implemented skimming

Bernanda Telalovic
1 parent 9314d941
Showing 9 changed files with 461 additions and 42 deletions Show diff stats
applyGenericFEI.py
configureReconstruction.py
example_configs/.fei.yaml.swp
example_configs/.reconstruction.yaml.swp
example_configs/fei.yaml
example_configs/skimming.yaml
load_configure_reconstruction.py
reconstructDecays.py
skimming.py
@@ -39,7 +39,8 @@ class FEI(object):
         self.pths = paths
         self.path = paths['main_path']
         self.verbose = verbose
-        self.fei_path = None
+        self.fei_path = []
+        self.skim_list = []
  
         add_aliases(aliases)
  
@@ -56,8 +57,14 @@ class FEI(object):
             database = self.fei['database']['db_location']
             database_txt = self.fei['database']['db_txt_location']
             prefix = self.fei['database']['prefix']
+
+            if 'central' in self.fei['database']['db_type']:
+
+                use_central_database(database_txt, LogLevel.WARNING,
+                        'fei_database')
+            else:
  
-            use_local_database(database_txt, database,
+                use_local_database(database_txt, database,
                            True, LogLevel.WARNING)
  
             b_cuts = build_cuts(self.fei['params']['B_extra_cut'])
@@ -123,6 +130,40 @@ class FEI(object):
  
  
  
+    def fei_skim(self, *args, **kwargs):
+
+        if self.fei.get('fei_skimming', False):
+        
+            self.skim_list = []
+            skimming = self.fei['fei_skimming'].get('particles', {})
+
+            for particle in skimming.keys():
+
+                for tag in skimming[particle].keys():
+
+                    list_name = particle + ':' + tag
+
+                    if list_name not in self.skim_list:
+
+                        self.skim_list.append(str(list_name))
+
+                    for module in skimming[particle][tag].keys():
+
+                        if str(module) in 'cuts':
+
+                            cuts = build_cuts(skimming[particle][tag]['cuts'])
+
+                            applyCuts(list_name,
+                                      cuts,
+                                      path = self.path)
+                        else:
+                            loose_mc_match = skimming[particle][tag]\
+                                        [module].get('looseMCMatch', False)
+                            
+                            self.path.add_module(module, list_name,
+                                looseMCMatching = loose_mc_match)
+            
+
  
     def apply_configuration(self, *args, **kwargs):
  
@@ -131,6 +172,7 @@ class FEI(object):
  
             self.load_FEI()
             self.filter_FEI()
+            self.fei_skim()
         else:
  
             message(self.verbose, 'FEI implementation has been disabled.' +\
 import basf2
 from basf2 import *
  
-from stdCharged import *
-from stdPi0s import *
 from ipython_tools import handler
 import numpy as np
 from matplotlib import pyplot as plt
@@ -24,11 +22,12 @@ from config import load_config
 from aux_functions import * 
 from reconstructDecays import ReconstructDecays
 from applyGenericFEI import FEI
+from skimming import Skimming
  
 from reconstruction import add_reconstruction, add_mdst_output
-from ipython_tools import handler
+from beamparameters import add_beamparameters
+
  
-#TODO: implement skimming and inporting/outputting udst files
  
 class ConfigureReconstruction(object):
  
@@ -123,6 +122,8 @@ class ConfigureReconstruction(object):
  
         if confs['reconstruction'] is not None:
             self.recon = load_config(confs['reconstruction'])
+            if not self.recon.get('implement', False):
+                self.recon = False
         else:
             message(verbose, 'No decay reconstruction config\
                     given. Will not preform signal side\
@@ -131,13 +132,24 @@ class ConfigureReconstruction(object):
  
         if confs['fei'] is not None:
             self.fei = load_config(confs['fei'])
+            if not self.fei.get('implement', False):
+                self.fei = False
         else:
             message(verbose, 'No FEI parameters given. \
                     Will not apply FEI.')
             self.fei = False
  
         if confs['skimming'] is not None:
+
             self.skimming = load_config(confs['skimming'])
+            self.truth_skim_path = create_path()
+            self.fei_skim_path = create_path()
+
+            self.fei_particle_list = []
+            self.truth_particle_list = []
+
+            if not self.skimming.get('implement', False):
+                self.skimming = False
         else:
             message(verbose, 'Skimming config not specified.')
             self.skimming = False
@@ -163,10 +175,23 @@ class ConfigureReconstruction(object):
         self.fei_tag = False
  
         if self.recon:
-            if self.recon.get('FEI_tag', False):
+            if self.recon.get('FEI_tag', False) and \
+                    self.recon.get('implement', False):
  
                 self.fei_tag = True
-                self.in_files['SIG'] = self.in_files['FEI']
+
+                if self.in_files.get('SIG', False) and \
+                        self.in_files.get('FEI', False):
+                
+                    self.in_files['SIG'] = self.in_files['FEI']
+
+        self.skim_fei_tag = False
+
+        if self.skimming:
+            if self.skimming.get('FEI_tag', False):
+
+                self.skim_fei_tag = True
+                self.in_files
  
  
  
@@ -174,12 +199,21 @@ class ConfigureReconstruction(object):
  
     def input_files(self, path, files, *args, **kwargs):
  
-        inputMdstList(files['type'],     # type of data to be run on
-                      files['file_list'],  # list of data filenames
-                      path,                # the analysis path
-                      skipNEvents = files.get('skip_n_events', 0)) # num evts
-                                                                   # to skip
+        file_type = files.get('file_type', 'mdst')
+        skip_evts = files.get('skip_n_events', 0)
+
+        if 'mdst' in file_type:
+            inputMdstList(files['type'],     # type of data to be run on
+                          files['file_list'],  # list of data filenames
+                          path,                # the analysis path
+                          skipNEvents = skip_evts) # num evts to skip
+        elif 'udst' in file_type:
+            inputMdstList(files['type'],     # type of data to be run on
+                          files['file_list'],  # list of data filenames
+                          path,                # the analysis path
+                          skipNEvents = skip_evts) # num evts to skip
  
+            
  
  
  
@@ -198,8 +232,62 @@ class ConfigureReconstruction(object):
            processed but will instead be appended to the signal path and the
            final file will be outputed after the signal side analysis.'''
  
+        if self.skimming:
+
+            confs = {'skimming': self.skimming}
+            paths = {'fei_main_path': self.fei_skim_path,
+                     'truth_main_path': self.truth_skim_path}
+            
+            if self.skimming.get('fei_skimming', False):
+                
+                self.input_files(self.fei_skim_path,
+                                 self.in_files['fei_skim'])
+
+            if self.skimming.get('truth_skimming', False) and \
+                    not self.skim_fei_tag:
+                
+                self.input_files(self.truth_skim_path,
+                                 self.in_files['truth_skim'])
+
+            skim = Skimming(confs, paths, verbose = self.verbose)
+            skim.apply_configuration()
+
+            if self.skimming.get('fei_skimming', False):
+                
+                self.fei_particle_list = skim.fei_skim_list
+
+                if self.skim_fei_tag:    
+                    self.truth_skim_path.add_path(self.fei_skim_path)
+                else:
+                
+                    self.build_output_files(self.fei_skim_path,
+                                            skim,
+                                            'fei_skimming',
+                                            self.fei_particle_list)
+
+                    summaryOfLists(self.fei_particle_list)
+                    
+                    self.do_analysis(self.fei_skim_path)
+
+
+            if self.skimming.get('truth_skimming', False):
+                
+                self.truth_particle_list = skim.truth_skim_list
+
+                self.build_output_files(self.truth_skim_path,
+                                        skim,
+                                        'truth_skimming',
+                                        self.truth_particle_list)
+
+                summaryOfLists(self.truth_particle_list)
+                
+                self.do_analysis(self.truth_skim_path)
+
+
         if self.fei and self.fei.get('implement', False):
  
+            #beam_params = add_beamparameters(self.fei_path, 'SuperKEKB')
+
             confs = {'fei': self.fei}
             paths = {'main_path': self.fei_path}
  
@@ -214,7 +302,7 @@ class ConfigureReconstruction(object):
                 self.sig_path.add_path(self.fei_path)
             else:
                 self.build_output_files(self.fei_path,
-                                    reconstruction, 'FEI')
+                                    reconstruction, 'FEI', [])
                 self.do_analysis(self.fei_path)
  
  
@@ -225,6 +313,8 @@ class ConfigureReconstruction(object):
  
             if not self.fei_tag:
                 self.input_files(self.sig_path, self.in_files['SIG'])
+                #beam_params = add_beamparameters(self.sig_path, 'SuperKEKB')
+
  
             reconstruction = ReconstructDecays(confs, paths,
                                     verbose = self.verbose)
@@ -232,7 +322,7 @@ class ConfigureReconstruction(object):
             reconstruction.apply_configuration()
  
             self.build_output_files(self.sig_path,
-                                    reconstruction, 'SIG')
+                                    reconstruction, 'SIG', [])
             self.do_analysis(self.sig_path)
  
  
@@ -240,7 +330,7 @@ class ConfigureReconstruction(object):
  
  
     def build_output_files(self, path, reconstruction,
-                           recon_tag, *args, **kwargs):
+                           recon_tag, lists, *args, **kwargs):
  
         '''Builds the output files according to the specifications in the
            configs. Will only build them if the output configs are given.
@@ -248,6 +338,47 @@ class ConfigureReconstruction(object):
            Params:
                @path: (B2Path) the analysis path that is to be '''
  
+        if 'skimming' in recon_tag:
+
+            self.build_udst_file(path, reconstruction, recon_tag, lists)
+
+        else:
+
+            self.build_mdst_file(path, reconstruction, recon_tag)
+
+
+
+
+    def build_udst_file(self, path, reconstruction, recon_tag, lists):
+
+        output = self.out_files.get(recon_tag, False)
+        output_file_tag = None
+
+        if output:
+            filetag = output['filetag']
+            file_path = output['path']
+
+            if file_path[-1] is not '/':
+                file_path += '/'
+
+            output_file_tag = file_path + filetag
+        else:
+            message(self.verbose, 'WARNING: No output filename given! Will not
+                    create output files!')
+
+        output_list = self.out_variables.get(recon_tag, [])
+
+        skimOutputUdst(self.skimming[recon_tag].get('decay_mode', recon_tag),
+                       skimParticleLists = lists,
+                       outputParticleLists = output_list,
+                       outputFile = output_file_tag)
+
+
+
+
+    def build_mdst_file(self, path, reconstruction,
+                        recon_tag, *args, **kwargs):
+
         output = self.out_files.get(recon_tag, False)
  
         if output:
@@ -6,6 +6,7 @@
  
 implement: True # must be turned on to do 
 database:
+    db_type: 'central'
     db_location: 
     db_txt_location:
     prefix:
@@ -0,0 +1,76 @@
+---
+# ---------------
+# Skimming Config
+# ---------------
+
+implement: True
+aliases:
+    sigProb: 'extraInfo(SignalProbability)'
+fei_skimming:
+    database:
+        db_type:
+        db_location:
+        db_txt_location:
+        prefix: 
+    KLong: False
+    chargedB: True
+    neutralB: False
+    semileptonic: False
+    hadronic: True
+    params:
+        B_extra_cut:
+            - 
+            - 
+        specific: False
+        convertedFromBelle: False
+        training: False
+        monitor: False
+        b2bii: False
+    particles:
+        <name>:
+            <coded_list_name>:
+                <module_to_add>:
+                    looseMCMatch: True
+                cuts:
+                    - 
+        <other_particle_name>:
+            <coded_list_name>:
+                cuts:
+                    - 
+
+truth_skimming:
+    extra_skim_list:
+        - 
+        - 
+    FSParticles:
+        type: loose
+
+    SIG:
+        Step0:
+            Upsilon(4S):
+                mu_cands:
+                    ch1: 'B-:sig_mu B+:generic'
+                    BuildRestOfEvent:
+                    cuts:
+                        - 
+        Step1:
+            B-:
+                sig_mu:
+                    ch1: 'D0:all mu-:loose'
+                    cuts:
+                        - 
+        Step2:
+            D0:
+                all:
+                    ch1: 'K-:loose pi+:loose'
+                    ch2: 'K-:loose pi+:loose pi0:veryLoose'
+                    ch3: 'K-:loose pi+:loose pi-:loose pi+:loose'
+                    ch4: 'pi-:loose pi+:loose '
+                    ch5: 'K_S0:all pi+:loose pi-:loose'
+                    ch6: 'K_S0:all pi+:loose pi-:loose pi0:veryLoose'
+                    ch7: 'K_S0:all pi0:veryLoose'
+                    ch8: 'K+:loose K-:loose'
+                    ch9: 'pi+:loose pi-:loose'
+                    ch10: 'K_S0:all K_S0:all'
+                    cut:
+                        - '0.86 < M < 3.0'
@@ -13,8 +13,8 @@ parser = argparse.ArgumentParser()
  
 parser.add_argument('--path',
                     dest = 'path',
-                    default = '/home/bernanda/honours/FEI/generic/' +\
-                               '1st_trial_run/configs/',
+                    default = '~/honours/FEI/generic/skims/phase2/' + \
+                            'MC9/charged/configs/run1',
                     type = str,
                     help = 'The directory of the config files (str).')
  
@@ -3,11 +3,14 @@ from basf2 import *
 from stdCharged import *
 from stdPi0s import *
 from stdPhotons import *
+from stdV0s import *
+
  
 from ipython_tools import handler
 import numpy as np
 from matplotlib import pyplot as plt
 import scipy as sp
+import re
  
 import ROOT
 from ROOT import Belle2
@@ -19,6 +22,7 @@ import os.path
 import sys
  
 from modularAnalysis import *
+from vertex import *
 import fei
 from fei import backward_compatibility_layer
  
@@ -64,6 +68,7 @@ class ReconstructDecays(object):
         self.verbose = verbose
         self.tags = {}
         self.particles = {}
+        self.skim_list = []
  
         add_aliases(aliases)
  
@@ -91,14 +96,18 @@ class ReconstructDecays(object):
             stdLooseMu(path = self.path)
             stdLooseK(path = self.path)
             stdPi0s(listtype = 'veryLoose', path = self.path)
+            stdPhotons(listtype = 'loose', path = self.path)
+            stdKshorts(path = self.path)
  
-        elif 'std' in typ:
+        elif 'std' in typ or 'tight' in typ:
  
             stdE(path = self.path)
             stdPi(path = self.path)
             stdMu(path = self.path)
             stdK(path = self.path)
             stdPi0(listtype = 'loose', path = self.path)
+            stdPhotons(listtype = 'tight', path = self.path)
+            stdKshorts(path = self.path)
  
         elif 'custom' is typ:
  
@@ -136,7 +145,7 @@ class ReconstructDecays(object):
  
         channel_list = list(particle.keys())
  
-        return [x for x in channel_list if x[:-1]=='ch']
+        return [x for x in channel_list if re.split(r'\d+', x)[0]=='ch']
  
  
  
@@ -157,6 +166,31 @@ class ReconstructDecays(object):
  
  
  
+    def vertex_fit(self, sig, step, particle, main_channel,
+                   channel, path, *args, **kwargs):
+
+        prt = sig[step][particle][main_channel]
+        vertexing = prt.get('fit_vertex', False)
+
+        if vertexing:
+
+            list_name = particle + ':' + main_channel + '_' + channel
+
+            decay_str = ''
+            if prt.get('decay_str_selection', False):
+
+                decay_str = list_name + '->' + sig[step][particle]\
+                                              [main_channel][channel]
+
+            vertexKFit(list_name,
+                       vertexing.get('acceptance_level', 0.0),
+                       decay_string = decay_str,
+                       constraint = vertexing.get('constraint', ''),
+                       path = self.path)
+
+
+
+
  
     def reconstruct_in_channels(self, sig, step, particle,
                                 main_channel, channels, *args, **kwargs):
@@ -171,15 +205,20 @@ class ReconstructDecays(object):
  
             cuts = build_cuts(sig[step][particle]\
                     [main_channel].get('cuts', []))
+
+            particle_and_tag = particle + ':' + main_channel + '_' + channel
  
-            decay_str = particle + ':' + main_channel + '_' + channel + ' -> '\
-                                + sig[step][particle][main_channel][channel] 
+            decay_str = particle_and_tag + ' -> ' +\
+                        sig[step][particle][main_channel][channel] 
  
             reconstructDecay(decay_str,
                              cuts,
-                             dmID = int(channel[-1]),
+                             dmID = int(channel.split('ch')[1]),
                              path = self.path)
  
+            self.vertex_fit(sig, step, particle,
+                            main_channel, channel, self.path)
+
  
  
  
@@ -213,7 +252,8 @@ class ReconstructDecays(object):
                 self._append_to_ls(self.tags[step][particle][main_channel],
                                    main_channel)
  
-                message(self.verbose, '        Main channel: ' + main_channel)
+                message(self.verbose, '        Main channel: ' +\
+                        main_channel)
  
                 particle_list_name = particle + ':' + main_channel
                 prt_list = [particle + ':' + main_channel + '_'\
@@ -238,39 +278,62 @@ class ReconstructDecays(object):
  
         '''Reconstructs decay given in reconstruction config'''
  
-        message(self.verbose, 'Preforming reconstruction...')
+        if self.cfgs['recon'].get('implement', False):
  
-        sig = self.cfgs['recon']['SIG']
+            message(self.verbose, 'Preforming reconstruction...')
+    
+            sig = self.cfgs['recon']['SIG']
+    
+            self.load_fs_particles()
+    
+            steps = self.get_steps()
+            n_steps = len(steps)
+    
+            for step_num in range(n_steps):
+    
+                step = 'Step' + str(n_steps - step_num - 1)
+                message(self.verbose, 'SIG RECON: ' + step)
+    
+                self.particles[step] = {}
+                self.tags[step] = {}
+    
+                self.reconstruct_in_particles(sig, step)
+
+            if self.cfgs['recon'].get('truth_skimming', False):
+
+                self.get_skim_list(sig)
  
-        self.load_fs_particles()
+        else:
+            message(self.verbose, 'Signal reconstruction disabled.' +\
+                    'Will not preform signal reconstruction')
  
-        steps = self.get_steps()
-        n_steps = len(steps)
  
-        for step_num in range(n_steps):
  
-            step = 'Step' + str(n_steps - step_num - 1)
-            message(self.verbose, 'SIG RECON: ' + step)
+    def get_skim_list(self, sig, *args, **kwargs):
  
-            self.particles[step] = {}
-            self.tags[step] = {}
+        for particle in sig.get('Step0', {}).keys():
  
-            self.reconstruct_in_particles(sig, step)
+            for list_name in sig['Step0'][particle].keys():
+            
+                particle_list_name = particle + ':' + list_name
  
+                self._append_to_ls(self.skim_list, particle_list_name)
  
  
+        extra_lists = self.cfgs['recon'].get('extra_skim_lists', False)
  
-    def apply_configuration(self):
+        if extra_lists:
  
-        add_aliases(self.cfgs['recon'].get('aliases', {}))
+            for ls in extra_lists:
+                self._append_to_ls(self.skim_list, ls)
  
-        if self.cfgs['recon'].get('implement', False):
  
-            self.reconstruct_decay()
  
-        else:
-            message(self.verbose, 'Signal reconstruction disabled.' +\
-                    'Will not preform signal reconstruction')
+    def apply_configuration(self):
+
+        add_aliases(self.cfgs['recon'].get('aliases', {}))
+
+        self.reconstruct_decay()
  
  
  
@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+
+from basf2 import *
+from modularAnalysis import *
+
+import sys
+sys.path.append('/home/bernanda/honours/scripts/')
+from applyGenericFEI import FEI
+from reconstructDecays import ReconstructDecays
+from aux_functions import *
+
+
+class Skimming(object):
+
+    def __init__(self, configs, paths, aliases = {},
+                 verbose = False, *args, **kwargs):
+
+        self.cfgs = configs
+        self.skimming = configs['skimming']
+        self.paths = paths
+        self.truth_path = paths['truth_main_path']
+        self.fei_path = paths['fei_main_path']
+        self.aliases = aliases
+        self.verbose = verbose
+
+        self.fei_skim_list = []
+        self.truth_skim_list = []
+
+        add_aliases(self.aliases)
+
+
+
+    def build_fei_config(self, *args, **kwargs):
+
+        config = {}
+        config['fei'] = {}
+
+        config['fei'] = self.skimming['fei_skimming']
+        config['fei']['implement'] = True
+
+        return config
+
+
+    def build_truth_config(self, *args, **kwargs):
+
+        config = {}
+        config['recon'] = {}
+
+        config['recon']['SIG'] = self.skimming['truth_skimming']
+        config['recon']['implement'] = True
+        config['recon']['truth_skimming'] = True
+
+        return config
+
+    def build_paths(self, path, *args, **kwargs):
+
+        config = {}
+        config = self.paths
+        config['main_path'] = path
+
+        return config
+
+
+    def skim(self, *args, **kwargs):
+
+        if self.skimming.get('implement', False):
+
+            if self.skimming.get('aliases', False):
+                add_aliases(self.skimming['aliases'])
+
+            if self.skimming.get('fei_skimming', False):
+    
+                message(self.verbose, 'Preforming FEI skimming.')
+    
+                configs = self.build_fei_config()
+                paths = self.build_paths(self.fei_path)
+    
+                fei_skimming = FEI(configs, paths,
+                                   aliases = self.aliases,
+                                   verbose = self.verbose)
+    
+                fei_skimming.apply_configuration()
+    
+                self.fei_skim_list = fei_skimming.skim_list
+    
+            if self.skimming.get('truth_skimming', False):
+    
+                message(self.verbose, 'Preforming truth skimming.')
+    
+                configs = self.build_truth_config()
+                paths = self.build_paths(self.truth_path)
+    
+                truth_skimming = ReconstructDecays(configs, paths,
+                                                aliases = self.aliases,
+                                                verbose = self.verbose)
+                truth_skimming.apply_configuration()
+    
+                self.truth_skim_list = truth_skimming.skim_list
+
+
+
+    def apply_configuration(self, *args, **kwargs):
+
+        self.skim()
+
+