mc_side_module.py

 
 
import ROOT
from ROOT import Belle2
 
import numpy as np
 
import tracking.validation.utilities as utilities
 
import tracking.harvest.refiners as refiners
import tracking.harvest.harvesting as harvesting
import tracking.harvest.peelers as peelers
ROOT.gSystem.Load("libtracking")
 
 
class MCSideTrackingValidationModule(harvesting.HarvestingModule):
    """Module to collect matching information about the found particles and to generate
       validation plots and figures of merit on the performance of track finding."""
 
    """ Expert level behavior:
        expert_level = default_expert_level: all figures and plots from this module except tree entries
        expert_level > default_expert_level: everything including tree entries
        expert_level <= default_expert_level//2: only basic figures
        default_expert_level//2 < expert_level < default_expert_level: basic figures and basic tree entries
    """
    
    default_expert_level = 10
 
    def __init__(
            self,
            name,
            contact,
            output_file_name=None,
            reco_tracks_name='RecoTracks',
            mc_reco_tracks_name='MCRecoTracks',
            expert_level=None):
        """Constructor"""
 
        output_file_name = output_file_name or name + 'TrackingValidation.root'
 
        super().__init__(foreach=mc_reco_tracks_name,
                         name=name,
                         output_file_name=output_file_name,
                         contact=contact,
                         expert_level=expert_level)
 
        
        self.reco_tracks_name = reco_tracks_name
 
        
        self.mc_reco_tracks_name = mc_reco_tracks_name
 
        
        self.track_match_look_up = None
 
        if self.expert_level >= self.default_expert_level:
            
            self.found_det_hit_ids = set()
 
            
            self.matched_det_hit_ids = set()
 
            
            self.clone_det_hit_ids = set()
 
            
            self.fake_det_hit_ids = set()
 
    def initialize(self):
        """Initialization signal at the start of the event processing"""
        super().initialize()
        self.track_match_look_up = Belle2.TrackMatchLookUp(self.mc_reco_tracks_name, self.reco_tracks_name)
 
    def prepare(self):
        """Collect some statistics about the pattern recognition tracks used for comparison to the MC tracks
 
        Executed once at the start of each event.
        """
        super().prepare()
        if self.expert_level >= self.default_expert_level:
            reco_tracks = Belle2.PyStoreArray(self.reco_tracks_name)
            track_match_look_up = self.track_match_look_up
 
            found_det_hit_ids = set()
            matched_det_hit_ids = set()
            clone_det_hit_ids = set()
            fake_det_hit_ids = set()
 
            for reco_track in reco_tracks:
                det_hit_ids = utilities.get_det_hit_ids(reco_track)
 
                found_det_hit_ids |= det_hit_ids
 
                if track_match_look_up.isAnyChargeMatchedPRRecoTrack(reco_track):
                    matched_det_hit_ids |= det_hit_ids
 
                if track_match_look_up.isAnyChargeClonePRRecoTrack(reco_track):
                    clone_det_hit_ids |= det_hit_ids
 
                if (track_match_look_up.isGhostPRRecoTrack(reco_track) or
                        track_match_look_up.isBackgroundPRRecoTrack(reco_track)):
                    fake_det_hit_ids |= det_hit_ids
 
            self.found_det_hit_ids = found_det_hit_ids
            self.matched_det_hit_ids = matched_det_hit_ids
            self.clone_det_hit_ids = clone_det_hit_ids
            self.fake_det_hit_ids = fake_det_hit_ids
 
    def pick(self, mc_reco_track):
        """Pick every MCRecoTrack"""
        return True
 
    def peel(self, mc_reco_track):
        """Looks at the individual Monte Carlo tracks and store information about them"""
        track_match_look_up = self.track_match_look_up
 
        # Analyse from the Monte Carlo reference side
        mc_reco_tracks = Belle2.PyStoreArray(self.foreach)
        multiplicity = mc_reco_tracks.getEntries()
 
        mc_particle = track_match_look_up.getRelatedMCParticle(mc_reco_track)
        is_primary = bool(mc_particle.hasStatus(Belle2.MCParticle.c_PrimaryParticle))
        mc_to_pr_match_info_crops = self.peel_mc_to_pr_match_info(mc_reco_track)
        mc_store_array_crops = peelers.peel_store_array_info(mc_reco_track, key="mc_{part_name}")
 
        crops = dict(is_primary=is_primary,
                     multiplicity=multiplicity,
                     **mc_to_pr_match_info_crops,
                     **mc_store_array_crops
                     )
 
        if self.expert_level >= self.default_expert_level:
            reco_track = track_match_look_up.getRelatedPRRecoTrack(mc_reco_track)
            mc_particle_crops = peelers.peel_mc_particle(mc_particle)
            hit_content_crops = peelers.peel_reco_track_hit_content(mc_reco_track)
 
            # Custom peel function to get single detector hit purities
            subdetector_hit_efficiency_crops = peelers.peel_subdetector_hit_efficiency(mc_reco_track, reco_track)
 
            mc_hit_efficiencies_in_all_pr_tracks_crops = self.peel_hit_efficiencies_in_all_pr_tracks(mc_reco_track)
 
            # Event Info
            event_meta_data = Belle2.PyStoreObj("EventMetaData")
            event_crops = peelers.peel_event_info(event_meta_data)
 
            # Store Array for easier joining
            store_array_crops = peelers.peel_store_array_info(reco_track, key="pr_{part_name}")
 
            # Information on PR reco track
            pr_purity_information = {
                "pr_hit_purity": track_match_look_up.getRelatedPurity(reco_track) if reco_track else float("nan"),
                **peelers.peel_subdetector_hit_purity(reco_track=reco_track, mc_reco_track=mc_reco_track,
                                                      key="pr_{part_name}")
            }
 
            # Information on TrackFinders
            trackfinder_crops = peelers.peel_trackfinder(reco_track)
 
            # Basic peel function to get Quality Indicators
            qualityindicator_crops = peelers.peel_quality_indicators(reco_track)
 
            crops.update(dict(**hit_content_crops,
                              **mc_particle_crops,
                              **subdetector_hit_efficiency_crops,
                              **mc_hit_efficiencies_in_all_pr_tracks_crops,
                              **event_crops,
                              **store_array_crops,
                              **pr_purity_information,
                              **trackfinder_crops,
                              **qualityindicator_crops
                              ))
 
        return crops
 
    def peel_mc_to_pr_match_info(self, mc_reco_track):
        """Extracts track-match information from the MCMatcherTracksModule results"""
        track_match_look_up = self.track_match_look_up
        return dict(
            is_matched=track_match_look_up.isAnyChargeMatchedMCRecoTrack(mc_reco_track),
            is_matched_correct_charge=track_match_look_up.isCorrectChargeMatchedMCRecoTrack(mc_reco_track),
            is_matched_wrong_charge=track_match_look_up.isWrongChargeMatchedMCRecoTrack(mc_reco_track),
            is_merged=track_match_look_up.isAnyChargeMergedMCRecoTrack(mc_reco_track),
            is_merged_correct_charge=track_match_look_up.isCorrectChargeMergedMCRecoTrack(mc_reco_track),
            is_merged_wrong_charge=track_match_look_up.isWrongChargeMergedMCRecoTrack(mc_reco_track),
            is_missing=track_match_look_up.isMissingMCRecoTrack(mc_reco_track),
            hit_efficiency=track_match_look_up.getRelatedEfficiency(mc_reco_track),
        )
 
    def peel_hit_efficiencies_in_all_pr_tracks(self, mc_reco_track):
        """Extracts hit efficiencies"""
        mc_det_hit_ids = utilities.get_det_hit_ids(mc_reco_track)
 
        hit_efficiency_in_all_found = utilities.calc_hit_efficiency(self.found_det_hit_ids,
                                                                    mc_det_hit_ids)
 
        unfound_hit_efficiency = 1.0 - hit_efficiency_in_all_found
 
        hit_efficiency_in_all_matched = utilities.calc_hit_efficiency(self.matched_det_hit_ids,
                                                                      mc_det_hit_ids)
 
        hit_efficiency_in_all_fake = utilities.calc_hit_efficiency(self.fake_det_hit_ids,
                                                                   mc_det_hit_ids)
 
        hit_efficiency_crops = dict(
            hit_efficiency_in_all_found=hit_efficiency_in_all_found,
            unfound_hit_efficiency=unfound_hit_efficiency,
            hit_efficiency_in_all_matched=hit_efficiency_in_all_matched,
            hit_efficiency_in_all_fake=hit_efficiency_in_all_fake,
        )
        return hit_efficiency_crops
 
    # Refiners to be executed on terminate #
    # #################################### #
 
    
    save_tree = refiners.save_tree(
        # using cond to suppress false doxygen warnings
        
        name="mc_tree",
        folder_name="mc_tree",
        above_expert_level=default_expert_level
        
    )
 
    
    save_tree_basic = refiners.save_tree(
        # using cond to suppress false doxygen warnings
        
        name="mc_tree",
        folder_name="mc_tree",
        above_expert_level=default_expert_level // 2,
        below_expert_level=default_expert_level
        
    )
 
    
    save_overview_figures_of_merit = refiners.save_fom(
        
        name="{module.id}_overview_figures_of_merit",
        title="Overview figures in {module.title}",
        aggregation=np.nanmean,
        key="{part_name}",
        select={"is_matched": "finding efficiency", "hit_efficiency": "hit efficiency", },
        filter_on="is_primary",
        description="""
finding efficiency - the ratio of matched primary Monte Carlo tracks to all Monte Carlo tracks
hit efficiency - the ratio of hits picked up by the matched pattern recognition track of primary Monte Carlo tracks
"""
        
    )
 
    # Default refiners that can be disabled with a lower expert_level
    # #################################### #
 
    
    save_hit_efficiency_histogram = refiners.save_histograms(
        
        above_expert_level=default_expert_level - 1,
        select={"hit_efficiency": "hit efficiency"},
        filter_on="is_primary",
        description="Not a serious plot yet.",
        
    )
 
    
    renaming_select_for_finding_efficiency_profiles = {
        'is_matched': 'finding efficiency',
        'd0_truth': 'd_{0}',
        'pt_truth': 'p_{t}',
        'multiplicity': 'multiplicity',
        'phi0_truth': '#phi',
    }
 
    
    save_finding_efficiency_profiles = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select=renaming_select_for_finding_efficiency_profiles,
        y='finding efficiency',
        y_binary=True,
        filter_on="is_primary",
        outlier_z_score=5.0,
        allow_discrete=True,
        
    )
 
    
    save_finding_efficiency_by_tan_lamba_profiles = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select={
            'is_matched': 'finding efficiency',
            'tan_lambda_truth': 'tan #lambda'
        },
        y='finding efficiency',
        y_binary=True,
        filter_on="is_primary",
        outlier_z_score=5.0,
        lower_bound=-1.73,
        upper_bound=3.27,
        
    )
 
    
    save_finding_efficiency_by_tan_lamba_in_pt_groups_profiles = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select={
            'is_matched': 'finding efficiency',
            'tan_lambda_truth': 'tan #lambda'
        },
        y='finding efficiency',
        y_binary=True,
        filter_on="is_primary",
        groupby=[("pt_truth", [0.070, 0.250, 0.600])],
        outlier_z_score=5.0,
        lower_bound=-1.73,
        upper_bound=3.27,
        
    )
 
    # Make profiles of the hit efficiencies versus various fit parameters
    
    renaming_select_for_hit_efficiency_profiles = {
        'hit_efficiency': 'hit efficiency',
        'd0_truth': 'd_{0}',
        'pt_truth': 'p_{t}',
        'multiplicity': 'multiplicity',
        'phi0_truth': '#phi',
    }
 
    
    save_hit_efficiency_profiles = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select=renaming_select_for_hit_efficiency_profiles,
        y='hit efficiency',
        y_binary=True,
        filter_on="is_primary",
        outlier_z_score=5.0,
        allow_discrete=True,
        
    )
 
    
    save_hit_efficiency_by_tan_lambda_profiles = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select={
            'hit_efficiency': 'hit efficiency',
            'tan_lambda_truth': 'tan #lambda',
        },
        y='hit efficiency',
        y_binary=True,
        filter_on="is_primary",
        outlier_z_score=5.0,
        lower_bound=-1.73,
        upper_bound=3.27,
        
    )
 
    
    save_finding_efficiency_by_pt_profiles_groupbyCharge = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select={
            'is_matched': 'finding efficiency',
            'pt_truth': 'p_{t}',
        },
        y='finding efficiency',
        y_binary=True,
        filter_on="is_primary",
        groupby=[("charge_truth", [0.])],
        outlier_z_score=5.0,
        allow_discrete=True,
        
    )
 
    
    save_finding_efficiency_by_tan_lambda_profiles_groupbyCharge = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select={
            'is_matched': 'finding efficiency',
            'tan_lambda_truth': 'tan #lambda'
        },
        y='finding efficiency',
        y_binary=True,
        filter_on="is_primary",
        groupby=[("charge_truth", [0.])],
        outlier_z_score=5.0,
        lower_bound=-1.73,
        upper_bound=3.27,
        
    )
 
    
    save_hit_efficiency_by_pt_profiles_groupbyCharge = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select={
            'hit_efficiency': 'hit efficiency',
            'pt_truth': 'p_{t}',
        },
        y='hit efficiency',
        y_binary=True,
        filter_on="is_primary",
        groupby=[("charge_truth", [0.])],
        outlier_z_score=5.0,
        allow_discrete=True,
        
    )
 
    
    save_hit_efficiency_by_tan_lambda_profiles_groupbyCharge = refiners.save_profiles(
        
        above_expert_level=default_expert_level - 1,
        select={
            'hit_efficiency': 'hit efficiency',
            'tan_lambda_truth': 'tan #lambda',
        },
        y='hit efficiency',
        y_binary=True,
        filter_on="is_primary",
        groupby=[("charge_truth", [0.])],
        outlier_z_score=5.0,
        lower_bound=-1.73,
        upper_bound=3.27,
        
    )
 
    
    save_hit_efficiency_in_all_found_hist = refiners.save_histograms(
        
        above_expert_level=default_expert_level - 1,
        # renaming quantity to name that is more suitable for display
        select=dict(hit_efficiency_in_all_found="total hit efficiency vs. all reconstructed tracks")
        
    )
 
    
    save_missing_mc_tracks_hit_efficiency_in_all_found_hist = refiners.save_histograms(
        
        above_expert_level=default_expert_level - 1,
        filter_on="is_missing",  # show only the efficiencies of missing mc tracks
        # renaming quantity to name that is more suitable for display
        select=dict(hit_efficiency_in_all_found="total hit efficiency in all reconstructed tracks for missing mc tracks")
        
    )