tracking-parametrisation-tuner/parameterisations/utils/preselection.py

import ROOT


def preselection(
    cuts: str = "",
    input_file: str = None,
    outfile_postfix: str = "selected",
    tree_name: str = "PrParameterisationData.PrMCDebugReconstructibleLong/Tuple",
) -> str:
    """Function that apply a selection to given data.

    Args:
        cuts (str, optional): String specifying the selection. Defaults to "".
        input_file (str, optional): Defaults to None.
        outfile_postfix (str, optional): Defaults to "selected".
        tree_name (str, optional): Defaults to "PrParameterisationData.PrMCDebugReconstructibleLong/Tuple".

    Returns:
        str: Path to the output file.
    """
    rdf = ROOT.RDataFrame(tree_name, input_file)
    rdf = rdf.Filter(cuts, "Selection")
    out_file = input_file.strip(".root") + f"_{outfile_postfix}.root"
    rdf.Snapshot("Selected", out_file)
    return out_file


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input-file",
        type=str,
        help="Path to the input file",
    )
    parser.add_argument(
        "--cuts",
        type=str,
        default="chi2_comb < 5 && pt > 10 && p > 1500 && p < 100000 && pid != 11",
        help="Cuts of the preselection",
    )
    args = parser.parse_args()
    preselection(**vars(args))