SpartyJetConfig.py
No OneTemporary
Actions

Size

12 KB

Subscribers

None

SpartyJetConfig.py
View Options

	# Copyright (c) 2010-11, Pierre-Antoine Delsart, Kurtis Geerlings, Joey Huston,
	# Brian Martin, and Christopher Vermilion
	#
	#----------------------------------------------------------------------
	# This file is part of SpartyJet.
	#
	# SpartyJet is free software; you can redistribute it and/or modify
	# it under the terms of the GNU General Public License as published by
	# the Free Software Foundation; either version 3 of the License, or
	# (at your option) any later version.
	#
	# SpartyJet is distributed in the hope that it will be useful,
	# but WITHOUT ANY WARRANTY; without even the implied warranty of
	# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
	# GNU General Public License for more details.
	#
	# You should have received a copy of the GNU General Public License
	# along with SpartyJet; if not, write to the Free Software
	# Foundation, Inc.:
	# 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
	#----------------------------------------------------------------------

	import os, os.path
	import logging
	m_log = logging.getLogger('SpartyLog')
	outStr = logging.StreamHandler()
	format = logging.Formatter("%(module)-18s %(levelname)-8s %(message)s")
	outStr.setFormatter(format)
	m_log.addHandler(outStr)

	from ROOT import gSystem

	try:
	sj_dir = os.environ['SPARTYJETDIR']
	sj_libdir = os.path.join(sj_dir, 'libs')
	except:
	m_log.error('Environment variable SPARTYJETDIR not set. You need to do source setup.sh in top directory')
	import sys
	sys.exit(1)

	# Load SpartyJet

	_sj_libs=[
	'libExternal.so',
	# 'libSpartyCore.so'
	'libJetCore.so',
	'libIO.so',
	'libJetTools.so',
	'libFastJetTools.so',
	'libExternalTools.so',
	'libEventShape.so',
	'libSpartyDisplay.so',
	'libUserPlugins.so'
	]
	for lib in _sj_libs:
	# Cmake puts libs here, try that first
	if os.path.exists(os.path.join(sj_libdir, 'Debug', lib)):
	gSystem.Load(os.path.join(sj_libdir, 'Debug', lib))
	elif os.path.exists(os.path.join(sj_libdir, lib)):
	gSystem.Load(os.path.join(sj_libdir, lib))
	else:
	print 'SpartyJetConfig WARNING missing',os.path.join(sj_libdir,lib)

	from ROOT import SpartyJet as SJ
	from ROOT import fastjet as fj

	# external tools not in the SpartyJet or fastjet namespaces
	from ROOT import Njettiness


	# print a banner
	def centeredLine(l):
	return '\|' + l.center(78) + '\|'

	def printBanner():
	print '#' + ('-'*78) + '#'
	print centeredLine('SpartyJet 4.0')
	print centeredLine('P.-A. Delsart, K. Geerlings, J. Huston, B. Martin, C. K. Vermilion')
	print centeredLine('http://projects.hepforge.org/spartyjet')
	print centeredLine(' ')
	print centeredLine('A package for jet finding and visualization')
	print centeredLine('Jet finding is performed using the FastJet package.')
	print centeredLine('(see below)')
	print centeredLine(' ')
	print centeredLine('Many externally authored tools are provided with SpartyJet;')
	print centeredLine('see COPYING for details.')
	print centeredLine(' ')
	print centeredLine('If you use SpartyJet for a publication, please cite')
	print centeredLine('Prog. Part. Nucl. Phys. (60) 484, 2008 [arXiv:0712.2447].')
	print '#' + ('-'*78) + '#'

	printBanner()

	def syncMessageLevel():
	'''Gets spartyjet global message level
	and sets logger to same level'''
	import ROOT
	SJ = ROOT.SpartyJet
	ld = {
	SJ.DEBUG :logging.DEBUG,
	SJ.INFO :logging.INFO,
	SJ.WARNING:logging.WARNING,
	SJ.ERROR :logging.ERROR
	}

	m_log.setLevel(ld[SJ.GlobalLevel()])

	"""
	TODO: Should these two helper functions be part of an InputMakerFactory class?
	"""

	def createNtupleInputMaker(filename, treename="", inputprefix="GUESS", Nvar="",
	momentumVars=(), inputType=None, masslessMode=False, inputsuffix=""):
	""" Create & configure a NtupleInputMaker class

	The function will try to guess eveything from the given filename
	TODO: Figure out input name in more cases.
	"""
	from ROOT import TFile, TTree
	from ROOT import SpartyJet as SJ

	syncMessageLevel()

	f = TFile(filename)
	if not f.IsOpen():
	m_log.error('Could not open ROOT file with name ', filename)
	return

	# ---------------------------------
	# get tree
	tree = None
	if treename == "":
	# use the first TTree found
	keys = [k.GetName() for k in f.GetListOfKeys( ) ]
	for k in keys:
	t = f.Get(k)
	if isinstance(t, TTree):
	tree = t
	treename = k
	break
	else:
	tree = f.Get(treename)
	if not bool(tree) :
	m_log.error("Couldn't find tree in ",filename)
	return
	# ---------------------------------

	branches = [b.GetName() for b in tree.GetListOfBranches( ) ]


	# ---------------------------------
	# Guess input prefix if not set
	if inputprefix=="GUESS":
	# branches, lower case
	for bn in branches:
	# we'll look for vars starting with 'input'
	if bn.lower().startswith('input') and '_' in bn:
	inputprefix = bn[:bn.find('_')] # get the part before '_'
	break
	# we'll look for vars without prefix'
	if bn.lower() == 'eta' or bn.lower() == 'px':
	inputprefix = ''
	break
	if inputprefix=="GUESS":
	m_log.error("Couldn't guess proper prefix for input variables")
	return
	else:
	m_log.info("Found prefix input = "+inputprefix)
	# ---------------------------------



	if inputprefix == '':
	# retrieve all variables starting with inputprefix
	branches = [ b for b in branches if '_' not in b ]
	# retrieve all vars from the branch name above : the XX part in bla_XX
	vars = dict( ( b.lower(),b ) for b in branches )
	else:
	if not inputprefix.endswith('_'): inputprefix += '_'
	# retrieve all variables starting with inputprefix
	branches = [ b for b in branches if b.startswith(inputprefix) ]
	# retrieve all vars from the branch name above : the XX part in bla_XX
	vars = dict( ( b[b.find('_')+1:].lower(),b[b.find('_')+1:] ) for b in branches )


	# ---------------------------------
	# Guess kinematic variables
	if momentumVars==():
	# try px,py,pz,e
	vars_set = set(vars.keys())
	if vars_set.issuperset(set(['px','py','pz','e'])):
	momentumVars = tuple( vars[k] for k in ('px','py','pz','e') )
	elif vars_set.issuperset(set(['eta','phi','pt','e'])):
	momentumVars = tuple( vars[k] for k in ('eta','phi','pt','e') )
	elif vars_set.issuperset(set(['eta','phi','p_t','e'])):
	momentumVars = tuple( vars[k] for k in ('eta','phi','p_t','e') )
	elif vars_set.issuperset(set(['eta','phi','pt','m'])):
	momentumVars = tuple( vars[k] for k in ('eta','phi','pt','m') )
	elif vars_set.issuperset(set(['eta','phi','p_t'])):
	momentumVars = tuple( vars[k] for k in ('eta','phi','p_t') )
	elif vars_set.issuperset(set(['eta','phi','pt'])):
	momentumVars = tuple( vars[k] for k in ('eta','phi','pt') )

	if momentumVars == ():
	m_log.error("Couldn't guess kinematic input variables")
	return
	else:
	m_log.info("Found kinematic input = "+str(momentumVars))
	# ---------------------------------
	# Guess input type
	if inputType == None:
	vtype = _branchType(tree.GetBranch(inputprefix+momentumVars[0]))
	momkey = (momentumVars[0]+momentumVars[-1]).lower()
	inputType = {
	'pxevector_double' : SJ.NtupleInputMaker.PxPyPzE_vector_double ,
	'pxevector_float' : SJ.NtupleInputMaker.PxPyPzE_vector_float,
	'pxearray_double' : SJ.NtupleInputMaker.PxPyPzE_array_double,
	'pxearray_float' : SJ.NtupleInputMaker.PxPyPzE_array_float,
	'etaevector_double' : SJ.NtupleInputMaker.EtaPhiPtE_vector_double,
	'etaevector_float' : SJ.NtupleInputMaker.EtaPhiPtE_vector_float,
	'etaearray_double' : SJ.NtupleInputMaker.EtaPhiPtE_array_double,
	'etaearray_float' : SJ.NtupleInputMaker.EtaPhiPtE_array_float,
	'etamvector_double' : SJ.NtupleInputMaker.EtaPhiPtM_vector_double,
	'etamvector_float' : SJ.NtupleInputMaker.EtaPhiPtM_vector_float,
	'etamarray_double' : SJ.NtupleInputMaker.EtaPhiPtM_array_double,
	'etamarray_float' : SJ.NtupleInputMaker.EtaPhiPtM_array_float,
	'etaptvector_double' : SJ.NtupleInputMaker.EtaPhiPt_vector_double,
	'etaptvector_float' : SJ.NtupleInputMaker.EtaPhiPt_vector_float,
	'etaptarray_double' : SJ.NtupleInputMaker.EtaPhiPt_array_double,
	'etaptarray_float' : SJ.NtupleInputMaker.EtaPhiPt_array_float
	}[momkey+vtype]
	m_log.info("Input variables type = "+vtype)

	# ---------------------------------
	# Guess variable N
	# only needed if array input
	if Nvar == "":
	for nName in ['n','num','nparticle']:
	if nName in vars:
	Nvar = vars[nName]
	if Nvar == "":
	if 'array' in vtype:
	m_log.error("Couldn't guess proper input_n variable, please define manually.")
	return
	else:
	m_log.info("Couldn't guess proper input_n variable, input_n will be read from vector size.")
	else:
	m_log.info("Found prefix input_n = "+inputprefix+Nvar)
	# ---------------------------------



	input = SJ.NtupleInputMaker(inputType)
	input.set_prefix(inputprefix)
	input.set_n_name(Nvar)
	input.set_variables(*momentumVars)
	input.setFileTree(filename, treename)
	input.set_name("InputJet")
	input.set_masslessMode(masslessMode) #

	return input

	def getInputMaker(filename):
	"""
	Guesses the appropriate InputMaker given a file's extension.

	For .root files, calls createNtupleInputMaker below. Other supported formats:
	.hep (StdHep)
	.hepmc (HepMC)
	.lhe (Les Houches Accord event format)
	.UW (simple text format)
	.dat (assumes file can be read by StdTextInput)
	"""
	from ROOT import SpartyJet as SJ
	extensionMap = {
	'hepmc' : SJ.HepMCInput,
	'lhe' : SJ.LheInput,
	'UW' : SJ.HuskyInput,
	'dat' : SJ.StdTextInput
	}

	try:
	extensionMap['hep'] = SJ.StdHepInput
	except:
	pass

	extension = filename.split('.')[-1] # everything after the last '.'
	if extension == 'root':
	return createNtupleInputMaker(filename)
	elif extension in extensionMap.keys():
	return extensionMap[extension](filename)
	elif extension == 'hep':
	m_log.error(filename+" cannot be loaded. StdHepInput not enabled. To enable, edit MakefileCommonOpt and then type 'make'")
	import sys
	sys.exit(1)
	else:
	m_log.error("Can't identify file type from extension "+extension+"!")
	import sys
	sys.exit(1)


	_keep_alive = []

	def addFinalCut(builder, ptCut ):
	finalPtCut = SJ.JetPtSelectorTool(ptCut)
	builder.add_jetTool(finalPtCut)
	_keep_alive.append(finalPtCut)




	# utils ...
	def _branchType(b):
	from ROOT import TBranchElement
	if isinstance(b, TBranchElement):
	if 'double' in b.GetTypeName():
	return 'vector_double'
	else:
	return 'vector_float'
	else: # assume it's a TBranch
	if '/F' in b.GetTitle( ):
	return 'array_float'
	else:
	return 'array_double'


	# Wrapper for functions that take vectors as arguments
	def stdVector(*args):
	'''Make a std::vector from a python list'''
	from ROOT import std
	floatType = False
	for i in args: floatType \|= isinstance(i,float)

	if floatType:
	vec = std.vector( float )()
	for i in args: vec.push_back(i)
	return vec
	else:
	vec = std.vector( int )()
	for i in args: vec.push_back(i)
	return vec

	def writeCurrentFile(rootFile):
	''' Save a copy of the current script, as a TText, in the current ROOT file
	nothing happens if a ROOT file is not currently open!
	'''
	import sys, os.path
	from ROOT import TString, TText, TFile
	filename = sys.argv[0]
	if (not os.path.exists(filename)):
	m_log.error("Could not open file "+filename+"!")
	return
	s = TString()
	t = TText()
	# Too fresh
	#with open(filename) as input:
	input = open(filename)
	if input:
	for line in input:
	s += line
	t.SetName(os.path.basename(filename))
	t.SetText(0,0,s.Data())
	f = TFile(rootFile, 'UPDATE')
	t.Write()
	f.Close()

File Metadata

Mime Type: text/x-python
Expires: Thu, Apr 24, 6:40 AM (1 d, 21 h)
Storage Engine: blob
Storage Format: Raw Data
Storage Handle: 4887813
Default Alt Text: SpartyJetConfig.py (12 KB)

SpartyJetConfig.pyNo OneTemporaryActions

SpartyJetConfig.pyView Options

File Metadata

Event Timeline

SpartyJetConfig.py
No OneTemporary
Actions

SpartyJetConfig.py
View Options