waveform_plots.py

Go to the documentation of this file.

r"""!
@file waveform_plots.py
@brief `main()` shows how to use
[load_waveforms()](#waveform_plots.load_waveforms) and
[load_waveforms_from_IceFinal()](#waveform_plots.load_waveforms_from_IceFinal).
@note Altair is needed to make the plot.
   ```
   pip3 install altair[save]
   ```
"""
 
#  Examples are in the main function at the bottom of the script
 
import ROOT
import numpy as np
import polars as pl
from pathlib import Path
from initialise import load_pueoEvent_Dataset
from scipy.signal import lfilter
 
 
 
 
 
def load_waveforms_old(dataset: ROOT.pueo.Dataset, antennas: str = None) -> pl.DataFrame():
    r"""! Loads time domain signals
    @ingroup lwf
 
    @param[in] dataset       The output of #initialise.load_pueoEvent_Dataset
    @param[in] antennas      (optional) **whitespace separated list** specifying
                             [AntNum](#antenna_attributes.ALL_PUEO_ANTENNA_NAMES)s.
    @retval    waveforms     The desired [channels](#antenna_attributes.read_global_channel_mapping)
                             and their waveforms [Volts]
 
@todo Currently main instrument only, since oftentimes LF is not simulated
 
*   `antenna`:
    *   Users can optionally provide a string containing a whitespace separated list of antenna
        indices (ie [AntNum](#antenna_attributes.read_global_channel_mapping))
    *   If no antennas is specified, all MI channels will be loaded into the DataFrame.
    *   An exception is raised if the list contains invalid `AntNum`s.
 
* Output schema:
    ```
    $ GlobalChan                           <u32>
    $ AntNum                              <enum>
    $ AntIdx                               <u32>
    $ Pol                                 <enum>
    $ AntType                             <enum>
    $ SurfNum                               <u8>
    $ SurfChanNum                           <u8>
    $ PhiSector                             <u8>
    $ Ring                                  <u8>
    $ waveforms (volts)       <array[f64, 1024]>
    $ step size (nanoseconds)              <f64>
    ```
    """
    
    from antenna_attributes import read_global_channel_mapping, ALL_PUEO_ANTENNA_NAMES
 
    if antennas is None:
        channel_mapping: pl.DataFrame = read_global_channel_mapping()
    else:  # if user provides a list of antennas, filter out unused antennas
        try:  # see if the user had provided a bad antenna index
            antennas = antennas.split()
            pl.Series(antennas, dtype=ALL_PUEO_ANTENNA_NAMES)
        except pl.exceptions.InvalidOperationError as exc:
            print("InvalidOperationError:", exc)
            print("\n\nPerhaps you included a non-existent antenna index?")
            print("aborting operation at function: load_waveforms()")
            exit(1)
 
        channel_mapping = (
            read_global_channel_mapping().filter(pl.col("AntNum").is_in(antennas))
        )
 
    # Filter out LF channels
    channel_mapping = channel_mapping.filter(pl.col("AntType") != "LF") 
 
    useful = dataset.useful()
 
    ''' 
    waveforms = channel_mapping.select(
        pl.all().exclude("Comment"),
 
        #return list, not np.array
        pl.col("GlobalChan").map_batches(
            lambda s: [np.asarray(useful.volts[int(i)], dtype=float) for i in s]
        ).alias("waveforms (volts)"),
 
        pl.col("GlobalChan").map_batches(
            lambda s: [float(useful.dt[int(i)]) for i in s]
        ).alias("step size (nanoseconds)")
    )
    '''
    
    waveforms = channel_mapping.select(
        pl.all().exclude("Comment"),
 
        pl.col("GlobalChan").map_batches(
            lambda s: np.array([useful.volts[i] for i in s])
        ).alias("waveforms (volts)"),
 
        pl.col("GlobalChan").map_batches(
            lambda s: np.array([useful.dt[i] for i in s])
        ).alias("step size (nanoseconds)")
    )
    
    return waveforms
 
def load_waveforms(dataset: ROOT.pueo.Dataset, antennas: str = None) -> pl.DataFrame:
    from antenna_attributes import read_global_channel_mapping, ALL_PUEO_ANTENNA_NAMES
    import numpy as np
    import polars as pl
 
    if antennas is None:
        channel_mapping: pl.DataFrame = read_global_channel_mapping()
    else:
        try:
            antennas = antennas.split()
            pl.Series(antennas, dtype=ALL_PUEO_ANTENNA_NAMES)
        except pl.exceptions.InvalidOperationError as exc:
            print("InvalidOperationError:", exc)
            print("\n\nPerhaps you included a non-existent antenna index?")
            print("aborting operation at function: load_waveforms()")
            exit(1)
 
        channel_mapping = read_global_channel_mapping().filter(pl.col("AntNum").is_in(antennas))
 
    channel_mapping = channel_mapping.filter(pl.col("AntType") != "LF")
    
    useful = dataset.useful()
 
    waveforms = channel_mapping.select(
        pl.all().exclude("Comment"),
 
        pl.col("GlobalChan").map_batches(
            lambda s: pl.Series(
                [np.asarray(list(useful.volts[int(i)]), dtype=np.float64) for i in s]
            )
        ).alias("waveforms (volts)"),
 
        pl.col("GlobalChan").map_batches(
            lambda s: pl.Series([float(useful.dt[int(i)]) for i in s])
        ).alias("step size (nanoseconds)")
    )
    '''
    waveforms = waveforms.with_columns(
    pl.col("waveforms (volts)")
      .map_batches(window_batches)
      .alias("waveforms (volts)")
    )
    '''
    return waveforms
 
def upsample_waveforms(waveforms: pl.DataFrame, upsample_factor: int) -> pl.DataFrame:
    r"""!Calls SciPy's
    [resample()](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.resample.html)
    to upsample the waveforms.
    @ingroup lwf
    @param[in] waveforms        The output of #load_waveforms
    @param[in] upsample_factor
    @retval    upsampled_waveforms
 
    *   The following columns are required in the input:
        -# `waveforms (volts)`
        -# `step size (nanoseconds)`
    *   No additional columns will be attached to the output,
        `waveforms (volts)` and `step size (nanoseconds)` will be upsampled and replaced.
 
    """
    from scipy.signal import resample
    signal_length = len(waveforms["waveforms (volts)"][0])
 
    upsampled_waveforms = waveforms.with_columns(
        pl.col("waveforms (volts)").map_batches(
            lambda wf: resample(wf, signal_length * upsample_factor, axis=1)
        ),
        pl.col("step size (nanoseconds)") / upsample_factor
    )
 
    return upsampled_waveforms
 
 
 
def window_batches(wfs: pl.Series) -> np.ndarray:
    WINDOW_front=60
    WINDOW_back=60
    out = []
    for wf in wfs:
        wf = np.asarray(wf)
        imax = np.argmax(abs(wf))
    
        tmp = np.zeros_like(wf)
        lo = max(imax - WINDOW_front, 0)
        hi = min(imax + WINDOW_back+1, wf.size)
        tmp[lo:hi] = wf[lo:hi]
 
        out.append(tmp)
 
    return np.stack(out)
 
 
def load_waveforms_from_root(filepath, ev=0):
    from antenna_attributes import read_global_channel_mapping
    import ROOT
    import os
 
    channel_mapping: pl.DataFrame = read_global_channel_mapping()
    channel_mapping = channel_mapping.filter(pl.col("AntType") == "LF")
    
    channel_mapping_H = channel_mapping.filter(pl.col("Pol") == "H")
    channel_mapping_V = channel_mapping.filter(pl.col("Pol") == "V")
 
    lib_dir = str(os.getenv('PUEO_UTIL_INSTALL_DIR'))+'/lib64'
 
    ROOT.gSystem.Load(lib_dir+"/libPueoSim.so")
    ROOT.gSystem.Load(lib_dir+"/libNiceMC.so")
 
 
    rootFile = ROOT.TFile.Open(filepath)
    passTree = rootFile.Get("passTree")
 
    passTree.GetEntry(ev)
 
 
    det_evt = passTree.detectorEvents[1]
    
    def load_y_H(s: pl.Series) -> np.ndarray:
        return np.stack([
            np.array(det_evt.waveformsH[int(i)-192].GetY(), copy=True)
            for i in s
        ])
 
    def load_y_V(s: pl.Series) -> np.ndarray:
        return np.stack([
            np.array(det_evt.waveformsV[int(i)-200].GetY(), copy=True)
            for i in s
        ])
    waveformsH = channel_mapping_H.select(
        pl.all().exclude("Comment"),
 
        pl.col("GlobalChan")
          .map_batches(load_y_H)
          .alias("waveforms (volts)"),
       pl.col("GlobalChan").map_batches(
            lambda s: np.array([1/3 for i in s])
        ).alias("step size (nanoseconds)")
    )
    waveformsV = channel_mapping_V.select(
        pl.all().exclude("Comment"),
 
        pl.col("GlobalChan")
          .map_batches(load_y_V)
          .alias("waveforms (volts)"),
       pl.col("GlobalChan").map_batches(
            lambda s: np.array([1/3 for i in s])
        ).alias("step size (nanoseconds)")
    )
 
    waveforms = pl.concat([waveformsH, waveformsV], how="vertical")
    
     
    waveforms = waveforms.with_columns(
    pl.col("waveforms (volts)")
      .map_batches(filter_batch)
      .alias("waveforms (volts)")
    )
    
    
    waveforms = waveforms.with_columns(
    pl.col("waveforms (volts)")
      .map_batches(window_batches)
      .alias("waveforms (volts)")
    )
    
    return waveforms
 
 
def filter_batch(s: pl.Series) -> np.ndarray:
    # s is shape (n_rows, 1024)
    taps = np.array([1, 1, 1,  1,  1,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  
        0,  0,  0,  0, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        0,  0,  0,  0,  0,  0,  1,  1,  1,  1,  1,  1,  1,  1,  1, 1,  1,  1,
        1,  1,  1,  0,  0,  0,  0,  0, -1, -1, -1, -1, -1, -2, -2, -2, -2, -2,
        -2, -1, -1, -1, -1,  0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  2,  2,  2,
        2,  2,  1,  1,  1,  0,  0, -1, -1, -2, -2, -2, -2, -2, -2, -2, -2, -2,
        -1, -1,  0,  1,  1,  2,  2,  2,  2,  2,  2,  2,  2,  1,  1, 0, -1, -2,
        -2, -2, -4, -4, -2, -2, -1,  0,  1,  2,  2,  2,  4,  2,  2,  1,  0, -1,
        -2, -2, -2, -2, -2, -1,  1,  2,  4,  4,  2, -1, -2, -2, -2,  0,  1,  1,  1])
    return np.array([lfilter(taps, 1.0, np.asarray(w, dtype=np.float64)) for w in s])
    
 
 
def plot_waveforms(waveforms: pl.DataFrame, plot_name: str) -> None:
    r"""!Uses Vega-Altair to make waveform plots.
    @ingroup lwf
    @param[in] waveforms    The output of any other functions in @ref load_waveforms
    @param[in] plot_name    Be sure to specify the file format (eg. `my_plot.png` or `my_plot.html`)
 
* The following columns are required in the input:
    *   `waveforms (volts)`
    *   `step size (nanoseconds)`
    *   `PhiSector`
    *   `Ring`
    *   `Pol`
 
* Example: Plot the vpol channels of three antennas using the last event in `run0/`.\n
  Place the figure in the `/tmp` directory.
```python
run_zero_data = load_pueoEvent_Dataset(pueo_mc_data=Path("/home/all_my_runs"), run_number=0)
run_zero_data.last()
wf = load_waveforms(dataset=run_zero_data, antennas="403 111 201")
plot_waveforms(wf.filter(pl.col("Pol") == "V"), plot_name="/tmp/waveforms.html")
```
  @note The actual file name will be changed to `/tmp/Vpol_waveforms.html`.
        Had we chosen not to filter out the horizontally polarized channels,
        there will be another plot: `/tmp/Hpol_waveforms.html`.
 
  \htmlonly
  <iframe src="example_three_antenna_waveforms_from_pueoEvent.html" width="950" height="1000"></iframe>
  \endhtmlonly
    """
    import altair as alt
 
    # this forces function to error out instead of failing silently
    waveforms.select("PhiSector", "Ring")
 
    time = np.zeros_like(waveforms["waveforms (volts)"])
    signal_length = time.shape[1]
 
    for i, dt in enumerate(waveforms["step size (nanoseconds)"]):
        time[i] = dt * np.arange(signal_length)
 
    waveforms_list = (
        waveforms.with_columns(pl.Series(time).alias("time (ns)"))
        .explode("waveforms (volts)", "time (ns)")
        .partition_by("Pol")
    )
 
    plot_path = Path(plot_name)
    destination = plot_path.parent
 
    # make the plots, let Altair take care of subplot positions automatically
    for frame in waveforms_list:
        pol = frame["Pol"][0]
 
        chart = (
            alt.Chart(frame).mark_line()
            .encode(
                x='time (ns):Q',
                y=alt.Y('waveforms (volts)').title("[Volts]"),
            )
            .facet(
                column=alt.Column("PhiSector:O", header=alt.Header(labelFontSize=16, titleFontSize=16)),
                row=alt.Row("Ring:O", header=alt.Header(labelFontSize=16, titleFontSize=16))
            )
            .properties(title=f"{pol}Pol Waveforms")
            .configure_title(fontSize=16)
        )
 
        pn = destination / f"{pol}pol_{plot_path.name}"
        chart.save(f"{str(pn)}")
 
 
def load_waveforms_from_IceFinal(filepath: str, event_number: [int] = 0,
                                 antennas: str = None, upsample_factor: int = None,
                                 want_plots: bool = False,
                                 plot_name: str = None) -> tuple[pl.DataFrame(), np.array(1024)]:
    r"""! Loads time domain signals from an IceFinal file;
    users can optionally provide a list of antenna indices.
    @ingroup lwf
 
    @param[in] filepath            path to `IceFinal_*_allTree.root`
    @param[in] event_number        (optional) an integer or a  **list** of **integers**, defaults to event 0
    @param[in] antennas            (optional) **whitespace separated string**  specifying the antenna names
    @param[in] upsample_factor     (optional)
    @param[in] want_plots          (optional) defaults to False
    @param[in] plot_name           (optional)
    @retval    waveforms_frame     waveforms recorded by the antennas [Volts]
    @retval    time_mesh           time mesh corresponding to the waveforms [seconds].
 
* The function will error out if the `passTree` of the IceFinal file (see `filepath`) is empty.
  (ie. no triggered events)
* If no antennas are specified, all waveforms will be loaded into the DataFrame,
  which presents them in [ascending](#antenna_attributes.ALL_PUEO_ANTENNA_NAMES) `AntNum`.
* Example:
  ```
  load_waveforms_from_IceFinal("IceFinal_420_allTree.root", antennas="403 111 201", event_number=[0, 1])
  ```
  ```
  ┌──────────────┬────────┬─────────────┬─────────────────────┬───────────────────────┐
  │ event number ┆ AntNum ┆ pueoSim idx ┆ VPol waveforms (V)  ┆ truth (signal origin) │
  │ ---          ┆ ---    ┆ ---         ┆ ---                 ┆ ---                   │
  │ i32          ┆ enum   ┆ u8          ┆ array[f64, 1024]    ┆ array[f64, 3]         │
  ╞══════════════╪════════╪═════════════╪═════════════════════╪═══════════════════════╡
  │ 0            ┆ 201    ┆ 1           ┆ […]                 ┆ […]                   │
  │ 0            ┆ 403    ┆ 11          ┆ […]                 ┆ […]                   │
  │ 0            ┆ 111    ┆ 40          ┆ […]                 ┆ […]                   │
  │ 1            ┆ 201    ┆ 1           ┆ […]                 ┆ […]                   │
  │ 1            ┆ 403    ┆ 11          ┆ […]                 ┆ […]                   │
  │ 1            ┆ 111    ┆ 40          ┆ […]                 ┆ […]                   │
  └──────────────┴────────┴─────────────┴─────────────────────┴───────────────────────┘
  ```
* The waveforms can be upsampled by @p upsample_factor.
* `truth (signal origin)` is a unit vector that represents the true direction of the incoming
  signal.
    * Note that in the IceFinal files, the variable `RFdir_payload` is the direction the signal is
      travelling **to**.
    * Therefore, to obtain the direction that the signal is coming **from**,
      we need the opposite vector: \n
      `signal_origin` = - `RFdir_payload`
    * This is useful for verifying the correctness of the [correlation maps](@ref example_corr_map).
 
* The function also optionally creates waveform plots and saves to the working directory
  (ie. wherever the user calls this script) via the internal function #plot_waveforms().
  Below is an example of a three-antenna plot (actually two plots would be made, one for each event):
  ```python
  load_waveforms_from_IceFinal(
          "IceFinal_420_allTree.root", antennas="403 111 201",
          want_plots=True, event_number=[0, 1])
  ```
  \anchor exampleThreeAntennaPlot
  \htmlonly
  <iframe src="example_three_antenna_waveforms.html" width="950" height="1000"></iframe>
  \endhtmlonly
  @todo **currently main instrument only**
 
  @warning Will abort if non-existent antenna indices are provided.
  ```
    load_waveforms_from_IceFinal("IceFinal_420_allTree.root", antennas="201 511")
  ```
  In this case, Antenna `511` does not exist, since there are only 4 antennas in each `PhiSector`.\n
  Thus, you would see
  > InvalidOperationError: conversion from `str` to `enum` failed in column '' for 1 out of 2 values: ["511"]
  >
  > Ensure that all values in the input column are present in the categories of the enum datatype.
  >
  > Perhaps you included a non-existent antenna index?
  > aborting operation at function: load_waveforms_from_IceFinal()
    """
    #from antenna_attributes import generate_dataframe_from_pueodata_qrh, ALL_PUEO_ANTENNA_NAMES
    from antenna_attributes import read_global_channel_mapping, ALL_PUEO_ANTENNA_NAMES
 
    if antennas is None:
        channel_mapping: pl.DataFrame = read_global_channel_mapping()
 
    # turn int into a list[int] so that it's iterable
    if isinstance(event_number, int):
        event_number = [event_number]
 
    # if user provides a list of antennas, filter out unused antennas
    if antennas is not None:
        # see if the user had provided a bad antenna index
        try:
            antennas = antennas.split()
            pl.Series(antennas, dtype=ALL_PUEO_ANTENNA_NAMES)
        except pl.exceptions.InvalidOperationError as exc:
            print("InvalidOperationError:", exc)
            print("\n\nPerhaps you included a non-existent antenna index?")
            print("aborting operation at function: load_waveforms_from_IceFinal()")
            exit(1)
 
        _index_conversion = (
            generate_dataframe_from_pueodata_qrh().filter(pl.col("AntNum").is_in(antennas))
        )
    else:
        _index_conversion: pl.DataFrame = generate_dataframe_from_pueodata_qrh()
 
    ROOT.gSystem.Load("libNiceMC.so")
    ROOT.gSystem.Load("libPueoSim.so")
 
    _rootFile: ROOT.TFile = ROOT.TFile.Open(filepath)
    _passTree: ROOT.TTree = _rootFile.Get('passTree')
 
    waveforms: [pl.DataFrame] = []
    for ev in event_number:
 
        if ev < 0 or ev >= _passTree.GetEntries():
            print(f"Invalid event index \"{ev}\" specified for {filepath}")
            print(f"There are only {_passTree.GetEntries()} events.")
            exit(1)
 
        _passTree.GetEvent(ev)
        _df: pl.DataFrame = (
            _index_conversion
            .select(
                pl.lit(ev).alias("event number"),
                pl.col("AntNum", "pueoSim idx"),
                pl.col("pueoSim idx").map_batches(
                    lambda s: np.array([
                        _passTree.detectorEvents[0].waveformsV[i].GetY()
                        for i in s
                    ])
                ).alias("VPol waveforms (V)"),
                pl.col("pueoSim idx").map_batches(
                    lambda s: np.array([
                        -(_passTree.detectorEvents[0].RFdir_payload)
                        for _ in s
                    ])
                ).alias("truth (signal origin)"),
            )
        )
        waveforms.append(_df)
    print(waveforms)
    # time mesh in seconds
    time_mesh = np.array(_passTree.detectorEvents[0].waveformsV[0].GetX())
 
    # turn the list of waveform dataframes into one dataframe and apply SciPy upsampling if needed
    waveforms = pl.concat(waveforms)
 
    # if want_plots:
    #     _plot_waveforms_frame(                  # plot in nanoseconds
    #         _explode_waveforms_frame(waveforms, time_mesh * 1e9), plot_name=plot_name
    #     )
 
    if upsample_factor is not None:
        from scipy.signal import resample
        upsampled_waveforms = waveforms.with_columns(
            pl.col("VPol waveforms (V)").map_batches(
                lambda s: resample(s, len(time_mesh) * upsample_factor, axis=1)
            )
        )
        _, upsampled_time_mesh = resample(waveforms["VPol waveforms (V)"][0],
                                          num=len(time_mesh) * upsample_factor, t=time_mesh)
 
        return upsampled_waveforms, upsampled_time_mesh
 
    return waveforms, time_mesh
 
 
 
 
if __name__ == "__main__":
    # Below we illustrate the basic features of the Dataset class,
    # and how to pass it to the analysis functions.
    # For more pueo::Dataset class methods, see its header file:
    #   https://github.com/PUEOCollaboration/pueoEvent/blob/main/src/pueo/Dataset.h
 
    # Use run zero (for example) to initialize the Dataset
    # Warning: make sure that the run is not empty (has events),
    # otherwise pueo::Dataset's constructor will crash
    myDataset: ROOT.pueo.Dataset = load_pueoEvent_Dataset(pueo_mc_data=Path("/users/PAS0654/yuchieh/geometry_bug"), run_number=1)
    print(f"The current run is {myDataset.getCurrRun()}")
 
    # suppose we now want to use another run, maybe run one. Warning: `/tmp/run1/` has to exist!
    #if not myDataset.loadRun(1, ROOT.pueo.Dataset.PUEO_MC_DATA):
    #    exit(1)
    #else:
    #    print(f"The current run is {myDataset.getCurrRun()}")
    print(f"The current run is {myDataset.getCurrRun()}")
 
    # Suppose we want to analyze the 15th entry
    if not myDataset.getEntry(1):
        print("no this entry")
        exit(1)
    else:
        print(f"Loaded event {myDataset.current()}")
 
    # We could also analyze the last event
    # myDataset.last()
 
    # ---------------------------- Examples ---------------------------- #
 
    # Example 1: plot the VPol waveforms of just three antennas:
    #wf: pl.DataFrame = load_waveforms(dataset=myDataset, antennas="403 111 201")
    #plot_waveforms(wf.filter(pl.col("Pol") == "V"), plot_name="./three_waveforms.png")
 
    # Example 2: plot upsampled VPol and HPol waveforms of all antennas
    #wf: pl.DataFrame = load_waveforms(dataset=myDataset)
    for ev in range(2):
        wf: pl.DataFrame = load_waveforms(dataset=myDataset, antennas="119 219 319 419 120 220 320 420")
        #wf: pl.DataFrame = load_waveforms_from_root("/users/PAS0654/yuchieh/geometry_bug/run0/IceFinal_0_allTree.root", ev=ev)
        #up: pl.DataFrame = upsample_waveforms(waveforms=wf, upsample_factor=3)
        plot_waveforms(wf, plot_name=f"/users/PAS0654/yuchieh/geometry_bug/GeoTest_MI_Evt{ev}_update.png")