diff --git a/examples/FCCee/import/AddAnalyzers.h b/examples/FCCee/import/AddAnalyzers.h new file mode 100644 index 0000000000..ba64f9a069 --- /dev/null +++ b/examples/FCCee/import/AddAnalyzers.h @@ -0,0 +1,10 @@ +#include "edm4hep/MCParticleData.h" + +ROOT::VecOps::RVec gen_particles() { + ROOT::VecOps::RVec result; + edm4hep::MCParticleData mcPart; + mcPart.momentum.x = 11; + result.push_back(mcPart); + + return result; +} diff --git a/examples/FCCee/import/test.py b/examples/FCCee/import/test.py new file mode 100644 index 0000000000..656c2e646b --- /dev/null +++ b/examples/FCCee/import/test.py @@ -0,0 +1,28 @@ +import FCCAnalyses +import ROOT + +ROOT.gROOT.SetBatch(True) + +def main(): + ''' + Example analysis entry point + ''' + + fccana = FCCAnalyses.Analysis('Test Analysis', 7) + + fccana.add_analyzers('examples/FCCee/import/AddAnalyzers.h') + + # fccana.add_files('examples/FCCee/import/test.root') + + dframe = fccana.get_dataframe() + dframe2 = dframe.Define("particles", "gen_particles()") + dframe3 = dframe2.Define("particles_pt", "MCParticle::get_pt(particles)") + hist = dframe3.Histo1D("particles_pt") + hist.Print() + + canvas = ROOT.TCanvas("canvas", "", 450, 450) + hist.Draw() + canvas.Print('test.pdf') + +if __name__ == '__main__': + main() diff --git a/python/FCCAnalyses.py b/python/FCCAnalyses.py new file mode 100644 index 0000000000..a8fc78ded9 --- /dev/null +++ b/python/FCCAnalyses.py @@ -0,0 +1,109 @@ +'''@package FCCAnalyses +Provides FCCAnalysis facilities as a Python import. +''' + +import os +import sys +import ROOT + +class Analysis: + ''' + The main class holding analysis settings + ''' + def __init__(self, name, datasource='local_files'): + self.name = name + self.description = '' + self.datasource = datasource + if isinstance(datasource, int): + self.n_events = datasource + self.datasource = 'n_events' + if datasource == 'local_files': + self.local_files = [] + + self.additional_analyzers = [] + + def set_description(self, description): + ''' + Set Analysis description. + ''' + self.description = description + + def get_dataframe(self): + ''' + Create and return ROOT RDataFrame + ''' + print('----> INFO: Loading standard FCCAnalyses analyzers...') + # TODO: find out which are actually needed + ROOT.gSystem.Load("libedm4hep") + ROOT.gSystem.Load("libpodio") + ROOT.gSystem.Load("libFCCAnalyses") + ROOT.gInterpreter.Declare("using namespace FCCAnalyses;") + _pod = ROOT.podio.ObjectID() + _edm = ROOT.edm4hep.ReconstructedParticleData() + _fcc = ROOT.dummyLoader() + + print('----> INFO: Loading additional analyzers from:') + for path in self.additional_analyzers: + print(' -', path) + + for path in self.additional_analyzers: + ROOT.gInterpreter.Declare(f'#include "{path}"') + + if self.datasource == 'local_files': + if not self.local_files: + print('----> ERROR: No input files provided. Aborting...') + sys.exit(3) + print('----> INFO: Analysis will run over the following input ' + 'files:') + for path in self.local_files: + print(' -', path) + + dframe = ROOT.RDataFrame('events', self.local_files) + + elif self.datasource == 'n_events': + print('----> INFO: Creating clear ROOT RDataFrame from ' + f'{self.n_events} events') + dframe = ROOT.RDataFrame(self.n_events) + + return dframe + + def add_files(self, infile_paths): + ''' + Adds file paths of local input files + ''' + if isinstance(infile_paths, str): + infile_paths = [infile_paths] + if isinstance(infile_paths, list): + for path in infile_paths: + if os.path.exists(path): + self.local_files.append(os.path.abspath(path)) + else: + print('----> ERROR: Input file not found!') + print(' ' + os.path.abspath(path)) + sys.exit(3) + else: + print('----> ERROR: Input files object is not a string nor a list') + sys.exit(3) + + def add_analyzers(self, add_analyzers_paths): + ''' + Loads additional analyzers to the ROOT gInterpreter. + The analyzers are JITed every time the analysis is run. + + Accepts either string or list of string. + ''' + if isinstance(add_analyzers_paths, str): + add_analyzers_paths = [add_analyzers_paths] + if isinstance(add_analyzers_paths, list): + for path in add_analyzers_paths: + if os.path.exists(path): + self.additional_analyzers.append(os.path.abspath(path)) + else: + print('----> ERROR: File with additional analyzers not ' + 'found!') + print(' ' + os.path.abspath(path)) + sys.exit(3) + else: + print('----> ERROR: Additional analyzer object is not a string nor ' + 'a list') + sys.exit(3)