From 0f9bfcc868f070027c87690a1d786a9c998a209a Mon Sep 17 00:00:00 2001 From: Vanya Belyaev Date: Tue, 20 Feb 2024 10:30:06 +0100 Subject: [PATCH] 1. `tree_reduce` : allow redefininition of existing variables (very useful for `tmva/chopping`) (only for 6.26<=ROOT --- ostap/frames/frames.py | 2 +- ostap/frames/tree_reduce.py | 17 +++++++++++------ 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/ostap/frames/frames.py b/ostap/frames/frames.py index 90feacbd..64f11422 100644 --- a/ostap/frames/frames.py +++ b/ostap/frames/frames.py @@ -117,7 +117,7 @@ def columns ( frame ) : """ names = [ str(c) for c in frame.GetColumnNames() ] if ( 6 , 16 ) <= root_info : - names += [ str(c) for c in frame.GetDefinedColumnNames() ] + names += [ str ( c ) for c in frame.GetDefinedColumnNames() ] return tuple ( sorted ( set ( names ) ) ) frame_columns = columns diff --git a/ostap/frames/tree_reduce.py b/ostap/frames/tree_reduce.py index ba605278..9d20e75d 100644 --- a/ostap/frames/tree_reduce.py +++ b/ostap/frames/tree_reduce.py @@ -8,7 +8,7 @@ # @author Vanya BELYAEV Ivan.Belyaev@itep.ru # @date 2018-06-16 # ============================================================================= -"""Helper module to ``reduce'' tree using frames +"""Helper module to `reduce' tree using frames - see Ostap.DataFrame - see ROOT.ROOT.RDataFrame """ @@ -34,7 +34,7 @@ if '__main__' == __name__ : logger = getLogger( 'ostap.frames.tree_reduce' ) else : logger = getLogger( __name__ ) # ============================================================================= -logger.debug ( "``Reduce'' TTree using ROOT::RDataFrame object") +logger.debug ( "`Reduce' TTree using ROOT::RDataFrame object") # ============================================================================= ## @class ReduceTree # Reduce TTree object using intermediate (temporary @@ -63,7 +63,7 @@ def __init__ ( self , tmp_keep = False , ## keep the temporary file silent = False ): ## silent processing - from ostap.frames.frames import DataFrame, frame_prescale + from ostap.frames.frames import DataFrame, frame_prescale, frame_columns frame = DataFrame ( chain ) report = None @@ -80,13 +80,18 @@ def __init__ ( self , ## add overall prescale (if requested) if 1 != prescale : frame = frame_prescale ( frame , prescale ) - + + avars = set ( frame_columns ( frame ) ) nvars = [] ## new variables for nv in new_vars : - frame = frame.Define ( nv , new_vars [ nv ] ) + if ( 6 , 26 ) <= root_info and nv in avars : + frame = frame.Redefine ( nv , new_vars [ nv ] ) ## REDEFINE!!! + else : + frame = frame.Define ( nv , new_vars [ nv ] ) nvars.append ( nv ) - + avars.add ( nv ) + from ostap.core.ostap_types import ( string_types , listlike_types , dictlike_types )