def GetCutList( self, cut, mathmod = math ): ''' This method allows to obtain a list with the events that satisfy the cuts given ''' cut, variables = FormatEvalExpr( cut, mathmod ) varstoadd = [ v for v in variables if v not in self.Variables ] if varstoadd: print 'WARNING: Loading additional variables to apply the cuts:', varstoadd self.BookVariables( *varstoadd ) values = [ self.Variables[ var ] for var in variables ] for ivar in xrange( len( variables ) ): cut = cut.replace( variables[ ivar ], 'values[ %i ][ ievt ]' %ivar ) return eval( '[ ievt for ievt in xrange( self.Nentries ) if ' + cut + ' ]' )
def GetVarEvents( self, *args, **kwargs ): ''' If an element in < args > is a variable, it gets the list of values for it. If it is an expression, it returns a list with the values corresponding to it. ''' cuts = kwargs.get( 'cuts', False ) mathmod = kwargs.get( 'mathmod', math ) if cuts: entries = self.GetCutList( cuts, mathmod ) else: entries = xrange( self.Nentries ) variables = [] trueargs = [] for v in args: if v in self.Variables: variables.append( v ) else: v, newv = FormatEvalExpr( v, mathmod ) variables += newv ''' The module is not imported, so the name must change ''' trueargs.append( v.replace( mathmod.__name__, 'mathmod' ) ) variables = list( set( variables ) ) variables.sort() variables.reverse() values = [ self.Variables[ var ] for var in variables ] nvars = len( variables ) output = [] for iv, arg in enumerate( trueargs ): for ivar in xrange( nvars ): arg = arg.replace( variables[ ivar ], 'values[ %i ][ ievt ]' %ivar ) output.append( eval( '[ %s for ievt in entries ]' %arg ) ) cmd = 'output[ 0 ]' for i in xrange( 1, len( output ) ): cmd += ', output[ %i ]' %i return eval( cmd )