-
Anthony Correia authoredAnthony Correia authored
persistence_csv.py 8.05 KiB
###############################################################################
# (c) Copyright 2019-2023 CERN for the benefit of the LHCb Collaboration #
# #
# This software is distributed under the terms of the GNU General Public #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING". #
# #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization #
# or submit itself to any jurisdiction. #
###############################################################################
"""Define configurable algorithms to dump hits and MC Particles into CSV file.
"""
from typing import Optional
from PyConf.tonic import configurable
from PyConf.application import make_data_with_FetchDataFromFile, make_odin
from RecoConf.hlt1_tracking import (
make_PrStoreUTHit_hits,
make_PrStoreSciFiHits_hits,
make_velo_full_clusters,
)
from .data_from_file import mc_unpackers
from .mc_checking import make_links_lhcbids_mcparticles_tracking_system
def erase_if_asked(path: str, erase: bool):
"""Erase the file in ``path`` if ``erase`` is set to ``True``."""
import os
import contextlib
if erase:
with contextlib.suppress(FileNotFoundError):
os.remove(path)
def csv_persistence_mchits(
odin_location=make_odin,
detector: Optional[str] = None,
output_path: str = "CSVDumper/mchits.csv",
erase: bool = True,
):
"""Dump the MC hits of a given detector into a CSV file.
Args:
detector: which detector's MC hits are stored. Set to ``None``
by default so ``detector`` must be given some values.
* ``VP``: velo MC hits
* ``UT``: UT MC hits
* ``FT``: scifi MC hits
output_path: path to the file where to save the CSV file
erase: whether to delete the existing CSV files, or just adding upon it
"""
from PyConf.Algorithms import PrCSVDumperMCHits
assert detector in ["VP", "UT", "FT"], "detector is " + str(detector)
mc_unpacked = mc_unpackers()
unpacked_mc_type = f"MC{detector}Hits"
erase_if_asked(output_path, erase)
print(f"MC hits of {detector} will be dumped to", output_path)
return PrCSVDumperMCHits(
MCHitLocation=mc_unpacked[unpacked_mc_type],
MCParticlesLocation=mc_unpacked["MCParticles"],
ODIN=odin_location(),
OutputPath=output_path,
)
@configurable
def csv_persistence_vp_mchits(output_path: str = "CSVDumper/vp_mchits.csv",
**kwargs):
"""Dump MC hits of the Velo."""
return csv_persistence_mchits(
detector="VP", output_path=output_path, **kwargs)
@configurable
def csv_persistence_ut_mchits(output_path: str = "CSVDumper/ut_mchits.csv",
**kwargs):
"""Dump MC hits of the UT."""
return csv_persistence_mchits(
detector="UT", output_path=output_path, **kwargs)
@configurable
def csv_persistence_ft_mchits(output_path: str = "CSVDumper/ft_mchits.csv",
**kwargs):
"""Dump MC hits of the SciFi."""
return csv_persistence_mchits(
detector="FT", output_path=output_path, **kwargs)
@configurable
def csv_persistence_event_info(
odin_location=make_odin,
output_path="CSVDumper/event_info.csv",
erase: bool = True,
):
from PyConf.Algorithms import PrCSVDumperEventInfo
erase_if_asked(output_path, erase)
print("Event info will be dumped to", output_path)
return PrCSVDumperEventInfo(
# VPLightClusterLocation=make_velo_light_clusters(),
VPClusterLocation=make_velo_full_clusters(),
FTHitsLocation=make_PrStoreSciFiHits_hits(),
UTHitsLocation=make_PrStoreUTHit_hits(),
ODIN=odin_location(),
OutputPath=output_path,
)
@configurable
def csv_persistence_mc_particles(
odin_location=make_odin,
output_path="CSVDumper/mc_particles.csv",
erase: bool = True,
extended: bool = False,
all_mc_particles: bool = True,
):
"""Dump the MC particles into a CSL file
Args:
odin_location: ODIN location
output_path: path to the CSV file where to write on
erase: whether to remove the CSV file it is exists
extended: whether to dump all the columns
all_mc_particles: whether to dump all the MC particles. Otherwise,
only the MC particles that have hits in the Velo, UT or SciFi
are dumped.
"""
from PyConf.Algorithms import PrCSVDumperMCParticles
links_to_lhcbids = make_links_lhcbids_mcparticles_tracking_system()
erase_if_asked(output_path, erase)
print("MC particles will be dumped to", output_path)
return PrCSVDumperMCParticles(
MCParticlesLocation=mc_unpackers()["MCParticles"],
ODIN=odin_location(),
MCTrackInfo=make_data_with_FetchDataFromFile("/Event/MC/TrackInfo"),
LinkerLocation=links_to_lhcbids,
OutputPath=output_path,
Extended=extended,
AllMCParticles=all_mc_particles,
)
@configurable
def csv_persistence_vp_hits(
odin_location=make_odin,
output_path="CSVDumper/hits_vp.csv",
erase: bool = True,
extended: bool = False,
):
"""Dump the Velo Hits into a CSV file
Args:
odin_location: ODIN location
output_path: path to the CSV file where to write on
erase: whether to remove the CSV file it is exists
extended: whether to dump all the columns
"""
from PyConf.Algorithms import PrCSVDumperVPHits
erase_if_asked(output_path, erase)
links_to_lhcbids = make_links_lhcbids_mcparticles_tracking_system()
print("VP hits will be dumped to", output_path)
return PrCSVDumperVPHits(
# VPLightClusterLocation=make_velo_light_clusters(),
VPClusterLocation=make_velo_full_clusters(),
MCParticlesLocation=mc_unpackers()["MCParticles"],
LinkerLocation=links_to_lhcbids,
ODIN=odin_location(),
OutputPath=output_path,
Extended=extended,
)
@configurable
def csv_persistence_ut_hits(
odin_location=make_odin,
output_path="CSVDumper/hits_ut.csv",
erase: bool = True,
extended: bool = False,
):
"""Dump the UT hits into a CSV file
Args:
odin_location: ODIN location
output_path: path to the CSV file where to write on
erase: whether to remove the CSV file it is exists
extended: whether to dump all the columns
"""
from PyConf.Algorithms import PrCSVDumperUTHits
erase_if_asked(output_path, erase)
links_to_lhcbids = make_links_lhcbids_mcparticles_tracking_system()
print("UT hits will be dumped to", output_path)
return PrCSVDumperUTHits(
UTHitsLocation=make_PrStoreUTHit_hits(),
MCParticlesLocation=mc_unpackers()["MCParticles"],
LinkerLocation=links_to_lhcbids,
ODIN=odin_location(),
OutputPath=output_path,
Extended=extended,
)
@configurable
def csv_persistence_ft_hits(
odin_location=make_odin,
output_path="CSVDumper/hits_ft.csv",
erase: bool = True,
extended: bool = False,
):
"""Dump the SciFi hits into a CSV file
Args:
odin_location: ODIN location
output_path: path to the CSV file where to write on
erase: whether to remove the CSV file it is exists
extended: whether to dump all the columns
"""
from PyConf.Algorithms import PrCSVDumperFTHits
erase_if_asked(output_path, erase)
links_to_lhcbids = make_links_lhcbids_mcparticles_tracking_system()
print("FT hits will be dumped to", output_path)
return PrCSVDumperFTHits(
FTHitsLocation=make_PrStoreSciFiHits_hits(),
MCParticlesLocation=mc_unpackers()["MCParticles"],
LinkerLocation=links_to_lhcbids,
ODIN=odin_location(),
OutputPath=output_path,
Extended=extended,
)