Source code for pipeline.Processing.modulewise_edges

"""This Module defines functions to build the true edges between the hits, use
a module from the origin vertex approach.
"""
import numpy as np
import pandas as pd


def get_modulewise_edges(hits: pd.DataFrame) -> np.ndarray:
    """Build the edges using a module-wise approach, from the production vertex.

    Args:
        hits: dataframe of hits for a given event, with columns ``vx``, ``vy``, ``vz``
            ``x``, ``y``, ``z``, ``particle_id``

    Returns:
        Array of all the edges, of shape (2, m), with ``m`` the number of edges.
    """
    signal = hits[  # Exclude noise etc
        ((~hits.particle_id.isna()) & (hits.particle_id != 0)) & (~hits.vx.isna())
    ]

    signal = signal.drop_duplicates(subset=["particle_id", "plane"])

    # Sort by increasing distance from production
    signal = signal.assign(
        R=np.sqrt(
            (signal.x - signal.vx) ** 2
            + (signal.y - signal.vy) ** 2
            + (signal.z - signal.vz) ** 2
        )
    )
    signal = signal.sort_values("R").reset_index(drop=False)

    # Handle re-indexing
    signal = signal.rename(columns={"index": "unsorted_index"}).reset_index(drop=False)
    signal.loc[
        signal["particle_id"] == 0, "particle_id"
    ] = np.nan  # Set 0 values to NaN

    # Group by particle ID
    signal_list = signal.groupby(["particle_id"], sort=False)["index"].agg(
        lambda x: list(x)
    )

    true_edges = []
    for row in signal_list.values:
        for i, j in zip(row[:-1], row[1:]):
            true_edges.append([i, j])

    true_edges = np.array(true_edges).T

    true_edges = signal.unsorted_index.values[true_edges]

    return true_edges