Source code for tools.xdigi2csvtools

"""Module that contains functions to configure the algorithms to execute,
to get the hits, MC particles and MC hits from a (X)DIGI file.
"""
import os.path as op
import typing
from contextlib import ExitStack
from definitions import dgen, dmoore
from Moore.config import Reconstruction




def get_algonames(
    detectors: typing.List[str],
    dump_event_info: bool = False,
    dump_mc_hits: bool = False,
) -> typing.List[str]:
    """Obtain the list of the names of the algorithms to run.

    Args:
        detectors: list of the detectors for which the hits and MC particles are dumped
        dump_event_info: whether to dump `event_info.csv`
        dump_mc_hits: wether to dump the MC hits

    Returns:
        List of the names of the algorithms to run
    """
    run_algonames = [f"hits_{detector}" for detector in detectors]
    run_algonames.append("mc_particles")
    if dump_mc_hits:
        run_algonames += [f"mchits_{detector}" for detector in detectors]
    if dump_event_info:
        run_algonames.append("event_info")
    return run_algonames





def get_algo_sequence(
    algonames: typing.List[str],
) -> typing.Callable[[], Reconstruction]:
    """A function that returns the list of algorithms to run to persist the hits
    and MC particles into CSV files, according to the configuration given as an
    input

    Args:
        algonames: list of the names of the algorithms to run

    Returns:
        Returns a function that returns the list of the algorithms,
        wrapped around a Moore ``Reconstruction`` object
    """

    def get_reconstruction():
        algos = [dmoore.algoname_to_algo[algoname]() for algoname in algonames]
        return Reconstruction("persistence_csv", algos)

    return get_reconstruction





def configure_algos(
    stack: ExitStack,
    algonames: typing.List[str],
    outdir: typing.Optional[str] = "persistence_csv",
    retina_clusters: bool = True,
    extended: bool = False,
    all_mc_particles: bool = False,
    erase: bool = True,
):
    """Configure the algorithms to run by changing the context
    (by basically doing a bunch of ``algo.bind(...)``)

    Args:
        stack: context manager. Modified in place
        algonames: list of the algorithms to run. We need this list in order only to
            to configure algorithms that will be run.
        retina_clusters: Whether to use Retina clusters
        extended: Whether to erase existing CSV files at the same locations
        outdir: path to the directory where the CSV files are saved
        erase: whether to erase an existing CSV file
    """
    from RecoConf.hlt1_tracking import make_velo_full_clusters
    from PyConf.Algorithms import VPRetinaFullClustering

    # Configuration the algorithms even though they are not run
    # Associates an algoname with the dictionnary of parameters to pass to the
    # algorithm to configure it
    algo_params = {  # initialisation
        algoname: {} for algoname in dmoore.algoname_to_algo.keys()
    }

    # Configure `erase` and `output_path` (all the algos have these variables)
    for algoname in dmoore.algoname_to_algo.keys():
        algo_params[algoname]["erase"] = erase
        algo_params[algoname]["output_path"] = op.join(outdir, f"{algoname}.csv")

    # Configure mc_particles
    algo_params["mc_particles"]["all_mc_particles"] = all_mc_particles
    algo_params["mc_particles"]["extended"] = extended

    # configure hits_{detector}
    for detector in dgen.all_detectors:
        algo_params[f"hits_{detector}"]["extended"] = extended

    with_velo = ("hits_velo" in algonames) or ("event_info" in algonames)

    for algoname, algo in dmoore.algoname_to_algo.items():
        if algoname in algonames:
            stack.enter_context(algo.bind(**algo_params[algoname]))

    # Supplementary configurations for retina clusters
    if with_velo and retina_clusters:
        stack.enter_context(
            make_velo_full_clusters.bind(
                make_full_cluster=VPRetinaFullClustering
            ),
        )