Basf2 specific examples¶
The following examples are not of interest to the general audience, but only for basf2 users.
Standard Simulation, Reconstruction and some nTuple Generation¶
import b2luigi as luigi
from b2luigi.basf2_helper import Basf2PathTask, Basf2nTupleMergeTask
from enum import Enum
import basf2
import modularAnalysis
import simulation
import generators
import reconstruction
from ROOT import Belle2
class SimulationType(Enum):
y4s = "Y(4S)"
continuum = "Continuum"
class SimulationTask(Basf2PathTask):
n_events = luigi.IntParameter()
event_type = luigi.EnumParameter(enum=SimulationType)
def create_path(self):
path = basf2.create_path()
modularAnalysis.setupEventInfo(self.n_events, path)
if self.event_type == SimulationType.y4s:
dec_file = Belle2.FileSystem.findFile('analysis/examples/tutorials/B2A101-Y4SEventGeneration.dec')
elif self.event_type == SimulationType.continuum:
dec_file = Belle2.FileSystem.findFile('analysis/examples/tutorials/B2A102-ccbarEventGeneration.dec')
else:
raise ValueError(f"Event type {self.event_type} is not valid. It should be either 'Y(4S)' or 'Continuum'!")
generators.add_evtgen_generator(path, 'signal', dec_file)
modularAnalysis.loadGearbox(path)
simulation.add_simulation(path)
path.add_module('RootOutput', outputFileName=self.get_output_file_name('simulation_full_output.root'))
return path
def output(self):
yield self.add_to_output("simulation_full_output.root")
@luigi.requires(SimulationTask)
class ReconstructionTask(Basf2PathTask):
def create_path(self):
path = basf2.create_path()
path.add_module('RootInput', inputFileNames=self.get_input_file_names("simulation_full_output.root"))
modularAnalysis.loadGearbox(path)
reconstruction.add_reconstruction(path)
modularAnalysis.outputMdst(self.get_output_file_name("reconstructed_output.root"), path=path)
return path
def output(self):
yield self.add_to_output("reconstructed_output.root")
@luigi.requires(ReconstructionTask)
class AnalysisTask(Basf2PathTask):
def create_path(self):
path = basf2.create_path()
modularAnalysis.inputMdstList('default', self.get_input_file_names("reconstructed_output.root"), path=path)
modularAnalysis.fillParticleLists([('K+', 'kaonID > 0.1'), ('pi+', 'pionID > 0.1')], path=path)
modularAnalysis.reconstructDecay('D0 -> K- pi+', '1.7 < M < 1.9', path=path)
modularAnalysis.fitVertex('D0', 0.1, path=path)
modularAnalysis.matchMCTruth('D0', path=path)
modularAnalysis.reconstructDecay('B- -> D0 pi-', '5.2 < Mbc < 5.3', path=path)
modularAnalysis.fitVertex('B+', 0.1, path=path)
modularAnalysis.matchMCTruth('B-', path=path)
modularAnalysis.variablesToNtuple('D0',
['M', 'p', 'E', 'useCMSFrame(p)', 'useCMSFrame(E)',
'daughter(0, kaonID)', 'daughter(1, pionID)', 'isSignal', 'mcErrors'],
filename=self.get_output_file_name("D_n_tuple.root"),
path=path)
modularAnalysis.variablesToNtuple('B-',
['Mbc', 'deltaE', 'isSignal', 'mcErrors', 'M'],
filename=self.get_output_file_name("B_n_tuple.root"),
path=path)
return path
def output(self):
yield self.add_to_output("D_n_tuple.root")
yield self.add_to_output("B_n_tuple.root")
class MasterTask(Basf2nTupleMergeTask):
n_events = luigi.IntParameter()
def requires(self):
for event_type in SimulationType:
yield self.clone(AnalysisTask, event_type=event_type)
if __name__ == "__main__":
luigi.process(MasterTask(n_events=1), workers=4)