Source code for pipeline.Evaluation.matching

import pandas as pd
import montetracko as mt


def perform_matching(
    df_tracks: pd.DataFrame,
    df_hits_particles: pd.DataFrame,
    df_particles: pd.DataFrame,
    min_track_length: int = 3,
    matching_fraction: float = 0.7,
    cure_clones: bool = False,
) -> mt.TrackEvaluator:
    """Perform matching and return the TrackEvaluator object for evaluation.

    Args:
        df_tracks: dataframe of tracks
        df_hits_particles: dataframe of hits-particles
        df_particles: dataframe of particles
        min_track_length: Minimum number of hits for a track to be kept
        matching_fration: Minimal matching fraction for the matching
        cure_clone: just a weird way of trying to remove clones that are matched
            to more than one particle

    Returns:
        TrackEvaluator object that contain the matched candidates.
    """
    trackEvaluator = mt.TrackMatcher(config={"fake_particle_id": 0}).full_matching(
        df_events_hits_particles=df_hits_particles,
        df_events_particles=df_particles,
        df_events_tracks_hits=df_tracks,
        matching_condition=(
            mt.matchcond.MinMatchingFraction(matching_fraction)
            & mt.matchcond.basics.BestMatchingParticles()
        ),
        track_condition=mt.matchcond.MinLengthTrack(min_track_length),
    )

    if cure_clones:
        df_candidates = trackEvaluator.dataframes["candidates"]
        n_matched_particles = (
            df_candidates.groupby(["event_id", "track_id"])["particle_id"]
            .count()
            .rename("n_matched_particles")
        )
        df_candidates = df_candidates.merge(
            n_matched_particles,
            how="left",
            on=["event_id", "track_id"],
        )
        df_clones = df_candidates[df_candidates["n_matched_particles"] >= 2].copy()
        df_clones["matching_fraction"] = (
            df_clones["n_matched_hits_particle_track"] / df_clones["n_hits_track"]
        )
        df_new_clones = df_clones[df_clones["matching_fraction"] > 0.85]
        df_missed_clones = pd.merge(
            df_clones,
            df_new_clones[["event_id", "track_id"]],
            how="outer",
            on=["event_id", "track_id"],
            indicator=True,
        )
        df_missed_clones = df_missed_clones[
            df_missed_clones["_merge"] == "left_only"
        ].drop("_merge", axis=1)
        trackEvaluator.dataframes["candidates"] = pd.concat(
            (
                df_candidates[df_candidates["n_matched_particles"] == 1],
                df_new_clones,
                df_missed_clones,
            ),
            axis=0,
        )

    return trackEvaluator