diff --git a/Train/dummy.py b/Train/dummy.py index a32264ae..80b97e1a 100644 --- a/Train/dummy.py +++ b/Train/dummy.py @@ -66,10 +66,10 @@ if not train.args.no_wandb: wandb.init( project=train.args.wandb_project, - config=wandb_config, + #config=wandb_config, ) wandb.save(sys.argv[0]) # Save python file - wandb.save(train.args.configFile) # Save config file + #wandb.save(train.args.configFile) # Save config file else: wandb.active=False diff --git a/Train/paper_trainer_Cocoa_simple.py b/Train/paper_trainer_Cocoa_simple.py new file mode 100644 index 00000000..d632934d --- /dev/null +++ b/Train/paper_trainer_Cocoa_simple.py @@ -0,0 +1,385 @@ +import sys +from argparse import ArgumentParser + +# import wandb +from DeepJetCore.wandb_interface import wandb_wrapper as wandb +import tensorflow as tf +from tensorflow.keras.layers import Concatenate, Dense, Dropout +from tensorflow.keras import Model + + +import training_base_hgcal +from Layers import LLRegulariseGravNetSpace, DummyLayer +from Layers import ProcessFeaturesCocoa +from Layers import RaggedGravNet +from Layers import PlotCoordinates +from Layers import TranslationInvariantMP +from Layers import LLExtendedObjectCondensationCocoa +from model_blocks import condition_input +from model_blocks import extent_coords_if_needed +from model_blocks import create_outputs +from callbacks import plotClusterSummary +from tensorflow.keras.layers import BatchNormalization + +#################################################################################################### +### Load Configuration ############################################################################# +#################################################################################################### + +parser = ArgumentParser('training') +parser.add_argument('--run_name', help="wandb run name", default="test") +parser.add_argument('--no_wandb', help="Don't use wandb", action='store_true') +parser.add_argument('--wandb_project', help="wandb_project", default="Paper_Models") +#parser.add_argument('--lr', help="learning rate", default="0.0001") + +# get the args so far but ignore the other args to pass them on +# this prevents errors with wandb in layers +pre_args,_ = parser.parse_known_args() + +if not pre_args.no_wandb: + wandb.init( + project=pre_args.wandb_project, + config={}, + ) + wandb.wandb().save(sys.argv[0]) # Save python file +else: + wandb.active=False + +#parses the rest of the arguments +train = training_base_hgcal.HGCalTraining(parser=parser) + + +PLOT_FREQUENCY = 6000 + +DENSE_ACTIVATION = 'elu' +DENSE_INIT = "he_normal" +d_shape = 64 +DENSE_REGULARIZER = tf.keras.regularizers.l2(1e-9) +NEIGHBOURS = [64,64,64,64] + +############################################################################### +### Define Model ############################################################## +############################################################################### + +def GravNet(name, + x, cprime, hit_energy, t_idx, + rs, + d_shape, + n_neighbours, + debug_outdir, + plot_debug_every, + space_reg_strength=-1e-2, + ): + + x= Dense(d_shape,activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x= Dense(d_shape,activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x= Dense(d_shape,activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + + x = Concatenate()([cprime, x]) + + x = BatchNormalization()(x) + + xgn, gncoords, gnnidx, gndist = RaggedGravNet( + name = "GravNet_"+name, + n_neighbours=n_neighbours, + n_dimensions=4, + n_filters=d_shape, + n_propagate=d_shape, + coord_initialiser_noise=1e-3, + feature_activation='tanh', + )([x, rs]) + + if space_reg_strength > 0: + gndist = LLRegulariseGravNetSpace(name=f'gravnet_coords_reg_{name}' , + record_metrics=True, + scale=space_reg_strength)([gndist, cprime, gnnidx]) + + gncoords = PlotCoordinates( + plot_every=plot_debug_every, + outdir = debug_outdir, + name=f'gncoords_{name}', + #publish='wandb' + )([gncoords, hit_energy, t_idx, rs]) + + x = DummyLayer()([x, gncoords]) #just so the branch is not optimised away, anyway used further down + x = BatchNormalization()(x) + + return Concatenate()([xgn, x]) + + +def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000): + """ + Function that defines the model to train + """ + + ########################################################################### + ### Pre-processing step ################################################### + ########################################################################### + + orig_input = td.interpretAllModelInputs(Inputs) + pre_processed = condition_input(orig_input, no_scaling=True, no_prime=False, new_prime=True) + + prime_coords = pre_processed['prime_coords'] + + c_coords = prime_coords + is_track = pre_processed['is_track'] + rs = pre_processed['row_splits'] + energy = pre_processed['rechit_energy'] + t_idx = pre_processed['t_idx'] + x = pre_processed['features'] + + x = ProcessFeaturesCocoa()(x) + + c_coords = PlotCoordinates( + plot_every=plot_debug_every, + outdir=debug_outdir, + name='input_c_coords', + # publish='wandb' + )([c_coords, energy, t_idx, rs]) + + c_coords = extent_coords_if_needed(c_coords, x, 3) + + x = Concatenate()([x, c_coords, is_track])#, tf.cast(t_idx, tf.float32)]) #JUST FOR TESTING FIXME + x = Dense(64, name='dense_pre_loop', activation=DENSE_ACTIVATION)(x) + + allfeat = [] + print("Available keys: ", pre_processed.keys()) + + ########################################################################### + ### Loop over GravNet Layers ############################################## + ########################################################################### + + + for i in range(len(NEIGHBOURS)): + + + x = GravNet(f'gncomb_{i}', x, prime_coords, energy, t_idx, rs, + d_shape, NEIGHBOURS[i], debug_outdir, plot_debug_every, space_reg_strength=-1e-2) + + #x = LayerNormalization()(x) + allfeat.append(x) + if len(allfeat) > 1: + x = Concatenate()(allfeat) + else: + x = allfeat[0] + + + ########################################################################### + ### Create output of model and define loss ################################ + ########################################################################### + + + x = Dense(128, + name=f"dense_final_{1}", + activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x = BatchNormalization()(x) + x = Dense(64, + name=f"dense_final_{2}", + activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x = Dense(64, + name=f"dense_final_{3}", + activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x = BatchNormalization()(x) + + pred_beta, pred_ccoords, pred_dist, \ + pred_energy, pred_energy_low_quantile, pred_energy_high_quantile, \ + pred_pos, pred_time, pred_time_unc, pred_id = \ + create_outputs(x, + n_ccoords=3, + n_classes=4, + n_pos = 3, + fix_distance_scale=True, + energy_factor=False, + is_track=is_track, + set_track_betas_to_one=False, + pred_e_factor=10) + + + pred_beta = LLExtendedObjectCondensationCocoa( + scale=1., + use_energy_weights=True, + record_metrics=True, + print_loss=train.args.no_wandb, + name="ExtendedOCLoss", + implementation = "atanh", + beta_loss_scale = 1.0, + too_much_beta_scale = 0.0, + energy_loss_weight = 1.0, + classification_loss_weight = 1.0, + position_loss_weight = 1.0, + timing_loss_weight = 0.0, + downweight_low_energy=False, + potential_scaling = 1.0, + train_energy_unc=False, + q_min = 0.1, + use_average_cc_pos = 0.9999)( + [pred_beta, pred_ccoords, pred_dist, pred_energy, pred_energy_low_quantile, + pred_energy_high_quantile, pred_pos, pred_time, pred_time_unc, pred_id, energy, + pre_processed['t_idx'] , pre_processed['t_energy'] , pre_processed['t_pos'] , + pre_processed['t_time'] , pre_processed['t_pid'] , pre_processed['t_spectator_weight'], + pre_processed['t_fully_contained'], pre_processed['t_rec_energy'], + pre_processed['t_is_unique'], pre_processed['row_splits']]) + + pred_ccoords = PlotCoordinates( + plot_every=plot_debug_every, + outdir = debug_outdir, + name='condensation', + #publish='wandb' + )([pred_ccoords, pred_beta, pre_processed['t_idx'], rs]) + + model_outputs = { + 'pred_beta': pred_beta, + 'pred_ccoords': pred_ccoords, + 'pred_energy': pred_energy, + #'pred_energy_low_quantile': pred_energy_low_quantile, + #'pred_energy_high_quantile': pred_energy_high_quantile, + 'pred_pos': pred_pos, + 'pred_time': pred_time, + 'pred_id': pred_id, + 'pred_dist': pred_dist, + 'rechit_energy': energy, + 'row_splits': pre_processed['row_splits'], + 't_idx': pre_processed['t_idx'], #JUST FOR TESTING FIXME + # 'no_noise_sel': pre_processed['no_noise_sel'], + # 'no_noise_rs': pre_processed['no_noise_rs'], + } + + # return DictModel(inputs=Inputs, outputs=model_outputs) + return Model(inputs=Inputs, outputs=model_outputs) + + +############################################################################### +### Set up training ########################################################### +############################################################################### + + +if not train.modelSet(): + train.setModel( + config_model, + td=train.train_data.dataclass(), + debug_outdir=train.outputDir+'/intplots', + plot_debug_every=PLOT_FREQUENCY, + ) + #train.applyFunctionToAllModels(RaggedGravNet.set_all_gn_space_trainable, False) #start with fixed GN space + train.setCustomOptimizer(tf.keras.optimizers.Adam())#clipnorm=1.)) + train.compileModel(learningrate=1e-4) + train.keras_model.summary() + train.saveModel(train.outputDir+'/model_init.h5') + +############################################################################### +### Callbacks ################################################################# +############################################################################### + + +samplepath = train.val_data.getSamplePath(train.val_data.samples[0]) +PUBLISHPATH = "" +PUBLISHPATH += [d for d in train.outputDir.split('/') if len(d)][-1] + + +cb = []#[NanSweeper()] #this takes a bit of time checking each batch but could be worth it + +cb += [ + plotClusterSummary( + outputfile=train.outputDir + "/clustering/", + samplefile=train.val_data.getSamplePath(train.val_data.samples[0]), + after_n_batches=500 + ) + ] + + +# cb += [wandbCallback()] + +############################################################################### +### Actual Training ########################################################### +############################################################################### + +train.change_learning_rate(1e-3) +train.trainModel( + nepochs=150, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 #average out more gradients + ) + + +#recompile +train.compileModel(learningrate=5e-4) +train.trainModel( + nepochs=400, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) + +#recompile +train.compileModel(learningrate=1e-4) +train.trainModel( + nepochs=600, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) + +# loop through model layers and turn batch normalisation to fixed +def fix_batchnorm(m): + for layer in m.layers: + if isinstance(layer, BatchNormalization): + layer.trainable = False + +#apply to all models +train.applyFunctionToAllModels(fix_batchnorm) + +#recompile +train.compileModel(learningrate=1e-3) +train.trainModel( + nepochs=650, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +#recompile +train.compileModel(learningrate=5e-4) +train.trainModel( + nepochs=800, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +train.compileModel(learningrate=1e-4) +train.trainModel( + nepochs=1200, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +train.compileModel(learningrate=1e-5) +train.trainModel( + nepochs=1400, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +train.compileModel(learningrate=1e-6) +train.trainModel( + nepochs=1500, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +try: + wandb.wandb().save(train.outputDir + "/KERAS_model.h5") +except: + print("Failed to save model") \ No newline at end of file diff --git a/models/simple_Cocoa/KERAS_model.h5 b/models/simple_Cocoa/KERAS_model.h5 new file mode 100644 index 00000000..ed7eea4c Binary files /dev/null and b/models/simple_Cocoa/KERAS_model.h5 differ diff --git a/models/simple_Cocoa/description.txt b/models/simple_Cocoa/description.txt new file mode 100644 index 00000000..883df757 --- /dev/null +++ b/models/simple_Cocoa/description.txt @@ -0,0 +1,5 @@ +Dataset: +/work/friemer/hgcalml/traindata_noEcut/dataCollection._n.djcdc (Cocoa Dataset with Eta Phi Cut) +->90% of energy of particle has to be within 0.5 eta and phi of mean position of hits + +Model trained for 1200 epochs (64h) \ No newline at end of file diff --git a/models/simple_Cocoa/paper_trainer_Cocoa_simple.py b/models/simple_Cocoa/paper_trainer_Cocoa_simple.py new file mode 100644 index 00000000..b9dd8fec --- /dev/null +++ b/models/simple_Cocoa/paper_trainer_Cocoa_simple.py @@ -0,0 +1,384 @@ +import sys +from argparse import ArgumentParser + +# import wandb +from DeepJetCore.wandb_interface import wandb_wrapper as wandb +import tensorflow as tf +from tensorflow.keras.layers import Concatenate, Dense, Dropout +from tensorflow.keras import Model + + +import training_base_hgcal +from Layers import LLRegulariseGravNetSpace, DummyLayer +from Layers import ProcessFeaturesCocoa +from Layers import RaggedGravNet +from Layers import PlotCoordinates +from Layers import TranslationInvariantMP +from Layers import LLExtendedObjectCondensation5 +from Layers import LLExtendedObjectCondensation3 +from model_blocks import condition_input +from model_blocks import extent_coords_if_needed +from model_blocks import create_outputs +from callbacks import plotClusterSummary +from tensorflow.keras.layers import BatchNormalization + +#################################################################################################### +### Load Configuration ############################################################################# +#################################################################################################### + +parser = ArgumentParser('training') +parser.add_argument('--run_name', help="wandb run name", default="test") +parser.add_argument('--no_wandb', help="Don't use wandb", action='store_true') +parser.add_argument('--wandb_project', help="wandb_project", default="Paper_Models") +#parser.add_argument('--lr', help="learning rate", default="0.0001") + +# get the args so far but ignore the other args to pass them on +# this prevents errors with wandb in layers +pre_args,_ = parser.parse_known_args() + +if not pre_args.no_wandb: + wandb.init( + project=pre_args.wandb_project, + config={}, + ) + wandb.wandb().save(sys.argv[0]) # Save python file +else: + wandb.active=False + +#parses the rest of the arguments +train = training_base_hgcal.HGCalTraining(parser=parser) + + +PLOT_FREQUENCY = 6000 + +DENSE_ACTIVATION = 'elu' +DENSE_INIT = "he_normal" +d_shape = 64 +DENSE_REGULARIZER = tf.keras.regularizers.l2(1e-9) +NEIGHBOURS = [64,64,64,64] + +############################################################################### +### Define Model ############################################################## +############################################################################### + +def GravNet(name, + x, cprime, hit_energy, t_idx, + rs, + d_shape, + n_neighbours, + debug_outdir, + plot_debug_every, + space_reg_strength=-1e-2, + ): + + x= Dense(d_shape,activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x= Dense(d_shape,activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x= Dense(d_shape,activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + + x = Concatenate()([cprime, x]) + + x = BatchNormalization()(x) + + xgn, gncoords, gnnidx, gndist = RaggedGravNet( + name = "GravNet_"+name, + n_neighbours=n_neighbours, + n_dimensions=4, + n_filters=d_shape, + n_propagate=d_shape, + coord_initialiser_noise=1e-3, + feature_activation='tanh', + )([x, rs]) + + if space_reg_strength > 0: + gndist = LLRegulariseGravNetSpace(name=f'gravnet_coords_reg_{name}' , + record_metrics=True, + scale=space_reg_strength)([gndist, cprime, gnnidx]) + + gncoords = PlotCoordinates( + plot_every=plot_debug_every, + outdir = debug_outdir, + name=f'gncoords_{name}', + #publish='wandb' + )([gncoords, hit_energy, t_idx, rs]) + + x = DummyLayer()([x, gncoords]) #just so the branch is not optimised away, anyway used further down + x = BatchNormalization()(x) + + return Concatenate()([xgn, x]) + + +def config_model(Inputs, td, debug_outdir=None, plot_debug_every=2000): + """ + Function that defines the model to train + """ + + ########################################################################### + ### Pre-processing step ################################################### + ########################################################################### + + orig_input = td.interpretAllModelInputs(Inputs) + pre_processed = condition_input(orig_input, no_scaling=True, no_prime=False, new_prime=True) + + prime_coords = pre_processed['prime_coords'] + + c_coords = prime_coords + is_track = pre_processed['is_track'] + rs = pre_processed['row_splits'] + energy = pre_processed['rechit_energy'] + t_idx = pre_processed['t_idx'] + x = pre_processed['features'] + + x = ProcessFeaturesCocoa()(x) + + c_coords = PlotCoordinates( + plot_every=plot_debug_every, + outdir=debug_outdir, + name='input_c_coords', + # publish='wandb' + )([c_coords, energy, t_idx, rs]) + + c_coords = extent_coords_if_needed(c_coords, x, 3) + + x = Concatenate()([x, c_coords, is_track])#, tf.cast(t_idx, tf.float32)]) #JUST FOR TESTING FIXME + x = Dense(64, name='dense_pre_loop', activation=DENSE_ACTIVATION)(x) + + allfeat = [] + print("Available keys: ", pre_processed.keys()) + + ########################################################################### + ### Loop over GravNet Layers ############################################## + ########################################################################### + + + for i in range(len(NEIGHBOURS)): + + + x = GravNet(f'gncomb_{i}', x, prime_coords, energy, t_idx, rs, + d_shape, NEIGHBOURS[i], debug_outdir, plot_debug_every, space_reg_strength=-1e-2) + + #x = LayerNormalization()(x) + allfeat.append(x) + if len(allfeat) > 1: + x = Concatenate()(allfeat) + else: + x = allfeat[0] + + + ########################################################################### + ### Create output of model and define loss ################################ + ########################################################################### + + + x = Dense(128, + name=f"dense_final_{1}", + activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x = BatchNormalization()(x) + x = Dense(64, + name=f"dense_final_{2}", + activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x = Dense(64, + name=f"dense_final_{3}", + activation=DENSE_ACTIVATION, kernel_initializer=DENSE_INIT, + kernel_regularizer=DENSE_REGULARIZER)(x) + x = BatchNormalization()(x) + + pred_beta, pred_ccoords, pred_dist, \ + pred_energy, pred_energy_low_quantile, pred_energy_high_quantile, \ + pred_pos, pred_time, pred_time_unc, pred_id = \ + create_outputs(x, + n_ccoords=3, + n_classes=4, + n_pos = 3, + fix_distance_scale=True, + energy_factor=False, + is_track=is_track, + set_track_betas_to_one=False, + pred_e_factor=10) + + + pred_beta = LLExtendedObjectCondensation5( + scale=1., + use_energy_weights=True, + record_metrics=True, + print_loss=train.args.no_wandb, + name="ExtendedOCLoss", + implementation = "atanh", + beta_loss_scale = 1.0, + too_much_beta_scale = 0.0, + energy_loss_weight = 1.0, + classification_loss_weight = 1.0, + position_loss_weight = 1.0, + timing_loss_weight = 0.0, + downweight_low_energy=False, + potential_scaling = 1.0, + train_energy_unc=False, + q_min = 0.1, + use_average_cc_pos = 0.9999)( + [pred_beta, pred_ccoords, pred_dist, pred_energy, pred_energy_low_quantile, + pred_energy_high_quantile, pred_pos, pred_time, pred_time_unc, pred_id, energy, + pre_processed['t_idx'] , pre_processed['t_energy'] , pre_processed['t_pos'] , + pre_processed['t_time'] , pre_processed['t_pid'] , pre_processed['t_spectator_weight'], + pre_processed['t_fully_contained'], pre_processed['t_rec_energy'], + pre_processed['t_is_unique'], pre_processed['row_splits']]) + + pred_ccoords = PlotCoordinates( + plot_every=plot_debug_every, + outdir = debug_outdir, + name='condensation', + #publish='wandb' + )([pred_ccoords, pred_beta, pre_processed['t_idx'], rs]) + + model_outputs = { + 'pred_beta': pred_beta, + 'pred_ccoords': pred_ccoords, + 'pred_energy': pred_energy, + #'pred_energy_low_quantile': pred_energy_low_quantile, + #'pred_energy_high_quantile': pred_energy_high_quantile, + 'pred_pos': pred_pos, + 'pred_time': pred_time, + 'pred_id': pred_id, + 'pred_dist': pred_dist, + 'rechit_energy': energy, + 'row_splits': pre_processed['row_splits'], + 't_idx': pre_processed['t_idx'], #JUST FOR TESTING FIXME + # 'no_noise_sel': pre_processed['no_noise_sel'], + # 'no_noise_rs': pre_processed['no_noise_rs'], + } + + # return DictModel(inputs=Inputs, outputs=model_outputs) + return Model(inputs=Inputs, outputs=model_outputs) + + +############################################################################### +### Set up training ########################################################### +############################################################################### + + +if not train.modelSet(): + train.setModel( + config_model, + td=train.train_data.dataclass(), + debug_outdir=train.outputDir+'/intplots', + plot_debug_every=PLOT_FREQUENCY, + ) + #train.applyFunctionToAllModels(RaggedGravNet.set_all_gn_space_trainable, False) #start with fixed GN space + train.setCustomOptimizer(tf.keras.optimizers.Adam())#clipnorm=1.)) + train.compileModel(learningrate=1e-4) + train.keras_model.summary() + train.saveModel(train.outputDir+'/model_init.h5') + +############################################################################### +### Callbacks ################################################################# +############################################################################### + + +samplepath = train.val_data.getSamplePath(train.val_data.samples[0]) +PUBLISHPATH = "" +PUBLISHPATH += [d for d in train.outputDir.split('/') if len(d)][-1] + + +cb = []#[NanSweeper()] #this takes a bit of time checking each batch but could be worth it + +cb += [ + plotClusterSummary( + outputfile=train.outputDir + "/clustering/", + samplefile=train.val_data.getSamplePath(train.val_data.samples[0]), + after_n_batches=500 + ) + ] + + +# cb += [wandbCallback()] + +############################################################################### +### Actual Training ########################################################### +############################################################################### + +train.change_learning_rate(1e-3) +train.trainModel( + nepochs=100, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 #average out more gradients + ) + + +#recompile +train.compileModel(learningrate=5e-4) +print('entering second training phase') +train.trainModel( + nepochs=350, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) + +#recompile +train.compileModel(learningrate=1e-4) +print('entering third training phase') +train.trainModel( + nepochs=500, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) + +# loop through model layers and turn batch normalisation to fixed +def fix_batchnorm(m): + for layer in m.layers: + if isinstance(layer, BatchNormalization): + layer.trainable = False + +#apply to all models +train.applyFunctionToAllModels(fix_batchnorm) + +#recompile +train.compileModel(learningrate=5e-4) +print('entering third training phase') +train.trainModel( + nepochs=750, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +#recompile +train.compileModel(learningrate=1e-4) +print('entering third training phase') +train.trainModel( + nepochs=1000, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +train.compileModel(learningrate=1e-5) +print('entering third training phase') +train.trainModel( + nepochs=1100, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +train.compileModel(learningrate=1e-6) +print('entering third training phase') +train.trainModel( + nepochs=1200, + batchsize=50000, + add_progbar=pre_args.no_wandb, + #additional_callbacks=cb, + collect_gradients = 4 + ) +try: + wandb.wandb().save(train.outputDir + "/KERAS_model.h5") +except: + print("Failed to save model") \ No newline at end of file diff --git a/modules/GravNetLayersRagged.py b/modules/GravNetLayersRagged.py index 12b06601..23bc5321 100644 --- a/modules/GravNetLayersRagged.py +++ b/modules/GravNetLayersRagged.py @@ -1601,7 +1601,6 @@ def __init__(self, 1.0, # recHitTime -> All zeros 1.0, # recHitHitR -> All zeros for tracks ]) - def get_config(self): config = {} @@ -1622,7 +1621,60 @@ def call(self, inputs): normalized = tf.where(is_track, normalized_tracks, normalized_hits) return normalized + +class ProcessFeaturesCocoa(ProcessFeatures): + def __init__(self, + **kwargs): + super().__init__(**kwargs) + self.mean_hit = tf.constant([ + 0.182, # recHitEnergy + 0.028, # recHitEta + 0.0, # recHitID -> don't normalize + 0.0074, # recHitTheta + 3012, # recHitR + 0.0, # recHitX -> centered around zero + 0.0, # recHitY -> centered around zero + 0.0, # recHitZ + 0.0, # recHitTime -> All zeros + 0.0, # recHitHitR + ]) + self.std_hit = tf.constant([ + 1.087, # recHitEnergy + 1.46, # recHitEta + 1.0, # recHitID -> don't normalize + 1.78, # recHitTheta + 829, # recHitR + 1096, # recHitX + 1096, # recHitY + 2604, # recHitZ + 1.0, # recHitTime -> All zeros + 1.0, # recHitHitR + ]) + self.mean_track = tf.constant([ + 15.173, # recHitEnergy + 0.0, # recHitEta + 0.0, # recHitID -> don't normalize + 1.57, # recHitTheta + 2574, # recHitR + 0.0, # recHitX -> centered around zero + 0.0, # recHitY -> centered around zero + 0.0, # recHitZ -> + 0.0, # recHitTime -> All zeros + 0.0, # recHitHitR -> All zeros for tracks + ]) + self.std_track = tf.constant([ + 19.126, # recHitEnergy + 1.31, # recHitEta + 1.0, # recHitID -> don't normalize + 0.93, # recHitTheta + 780, # recHitR + 990, # recHitX + 991, # recHitY + 2297, # recHitZ -> + 1.0, # recHitTime -> All zeros + 1.0, # recHitHitR -> All zeros for tracks + ]) diff --git a/modules/Layers.py b/modules/Layers.py index 54c0e1ff..824047c7 100644 --- a/modules/Layers.py +++ b/modules/Layers.py @@ -144,6 +144,9 @@ from GravNetLayersRagged import ProcessFeatures global_layers_list['ProcessFeatures']=ProcessFeatures +from GravNetLayersRagged import ProcessFeaturesCocoa +global_layers_list['ProcessFeaturesCocoa']=ProcessFeaturesCocoa + from GravNetLayersRagged import ManualCoordTransform global_layers_list['ManualCoordTransform']=ManualCoordTransform @@ -312,6 +315,7 @@ from LossLayers import LLFullObjectCondensationUncertainty, LLFullObjectCondensationID from LossLayers import LLExtendedObjectCondensation, LLExtendedObjectCondensation2 from LossLayers import LLExtendedObjectCondensation3, LLExtendedObjectCondensation4, LLExtendedObjectCondensation5 +from LossLayers import LLExtendedObjectCondensationCocoa from LossLayers import LLEdgeClassifier, AmbiguousTruthToNoiseSpectator, LLGoodNeighbourHood, LLKnnPushPullObjectCondensation from LossLayers import LLEnergySums,LLKnnSimpleObjectCondensation, LLPushTracks, LLFullOCThresholds, LLLocalEnergyConservation from LossLayers import LLRegulariseGravNetSpace, LLSpectatorPenalty @@ -355,6 +359,7 @@ global_layers_list['LLExtendedObjectCondensation3']=LLExtendedObjectCondensation3 global_layers_list['LLExtendedObjectCondensation4']=LLExtendedObjectCondensation4 global_layers_list['LLExtendedObjectCondensation5']=LLExtendedObjectCondensation5 +global_layers_list['LLExtendedObjectCondensationCocoa']=LLExtendedObjectCondensationCocoa global_layers_list['LLFullObjectCondensationID']=LLFullObjectCondensationID global_layers_list['LLGraphCondOCLoss']=LLGraphCondOCLoss global_layers_list['LLPFCondensates']=LLPFCondensates diff --git a/modules/LossLayers.py b/modules/LossLayers.py index 37a08dd4..a73c84fd 100644 --- a/modules/LossLayers.py +++ b/modules/LossLayers.py @@ -2178,6 +2178,7 @@ def __init__(self, *, energy_loss_weight=1., beta_push=0., implementation = '', global_weight=False, + train_energy_unc=True, **kwargs): """ Read carefully before changing parameters @@ -2231,7 +2232,7 @@ def __init__(self, *, energy_loss_weight=1., raise ValueError("huber_energy_scale>0 and alt_energy_loss exclude each other") - from object_condensation import Basic_OC_per_sample, PushPull_OC_per_sample, PreCond_OC_per_sample + from object_condensation import Basic_OC_per_sample, PushPull_OC_per_sample, PreCond_OC_per_sample, Hinge_OC_per_sample_atanhweighing from object_condensation import Hinge_OC_per_sample_damped, Hinge_OC_per_sample, Hinge_Manhatten_OC_per_sample from object_condensation import Dead_Zone_Hinge_OC_per_sample,Hinge_OC_per_sample_learnable_qmin, Hinge_OC_per_sample_learnable_qmin_betascale_position impl = Basic_OC_per_sample @@ -2255,6 +2256,8 @@ def __init__(self, *, energy_loss_weight=1., impl = Hinge_OC_per_sample_learnable_qmin if implementation == 'hinge_qmin_betascale_pos': impl = Hinge_OC_per_sample_learnable_qmin_betascale_position + if implementation == 'atanh': + impl = Hinge_OC_per_sample_atanhweighing self.implementation = implementation @@ -2318,6 +2321,7 @@ def __init__(self, *, energy_loss_weight=1., self.loc_time=time.time() self.call_count=0 self.beta_push = beta_push + self.train_energy_unc=train_energy_unc assert kalpha_damping_strength >= 0. and kalpha_damping_strength <= 1. @@ -2754,7 +2758,8 @@ def get_config(self): 'div_repulsion' : self.div_repulsion, 'dynamic_payload_scaling_onset': self.dynamic_payload_scaling_onset, 'beta_push': self.beta_push, - 'implementation': self.implementation + 'implementation': self.implementation, + 'train_energy_unc': self.train_energy_unc } base_config = super(LLFullObjectCondensation, self).get_config() return dict(list(base_config.items()) + list(config.items())) @@ -3053,6 +3058,10 @@ def calc_energy_correction_factor_loss(self, epred = pred_energy * t_dep_energies sigma = pred_uncertainty_high * t_dep_energies + 1.0 + + if not self.train_energy_unc: + sigma = tf.ones_like(sigma) + # Uncertainty 'sigma' must minimize this term: # ln(2*pi*sigma^2) + (E_true - E-pred)^2/sigma^2 @@ -3060,7 +3069,9 @@ def calc_energy_correction_factor_loss(self, prediction_loss = tf.math.divide_no_nan((t_energy - epred)**2, sigma**2) uncertainty_loss = tf.math.log(sigma**2) - + if not self.train_energy_unc: + uncertainty_loss = pred_uncertainty_high**2 + matching_loss = tf.debugging.check_numerics(matching_loss, "matching_loss") prediction_loss = tf.debugging.check_numerics(prediction_loss, "matching_loss") uncertainty_loss = tf.debugging.check_numerics(uncertainty_loss, "matching_loss") @@ -3278,6 +3289,112 @@ def calc_energy_correction_factor_loss(self, else: return prediction_loss, uncertainty_loss + matching_loss +class LLExtendedObjectCondensationCocoa(LLExtendedObjectCondensation): + ''' + Same as `LLExtendedObjecCondensation3` but uses total energy instead of correction as prediction, + position loss for [cos(phi), sin(phi), eta] and classification loss into categories for COCOA + ''' + + + def __init__(self, *args, **kwargs): + super(LLExtendedObjectCondensationCocoa, self).__init__(*args, **kwargs) + + + def calc_energy_correction_factor_loss(self, + t_energy, t_dep_energies, + pred_energy, pred_uncertainty_low, pred_uncertainty_high, + return_concat=False): + """ + Wrong name for parity, but calculates the total energy instead of the correction factor. + * t_energy -> Truth energy of shower + * t_dep_energies -> Sum of deposited energy IF clustered perfectly + * pred_energy -> predicted energy + * pred_uncertainty_low -> predicted uncertainty + * pred_uncertainty_high -> predicted uncertainty (should be equal to ...low) + """ + t_energy = tf.clip_by_value(t_energy,0.,1e12) + + prediction_loss = tf.math.log((t_energy - pred_energy)**2+1) + #prediction_loss = ((t_energy - pred_energy)/t_energy)**2 + + uncertainty_loss = pred_uncertainty_high**2 + pred_uncertainty_low**2 + + prediction_loss = tf.debugging.check_numerics(prediction_loss, "matching_loss") + uncertainty_loss = tf.debugging.check_numerics(uncertainty_loss, "matching_loss") + uncertainty_loss = tf.clip_by_value(uncertainty_loss, 0, 10) + + if return_concat: + return tf.concat([prediction_loss, uncertainty_loss], axis=-1) + else: + return prediction_loss, uncertainty_loss + + def calc_position_loss(self, t_pos, pred_pos): + if tf.shape(pred_pos)[-1] == 2: + raise ValueError("Position loss is implemented for cosphi, sinphi, eta") + if not self.position_loss_weight: + return 0.*tf.reduce_sum((pred_pos-t_pos)**2,axis=-1, keepdims=True) + + ploss = huber(tf.sqrt(tf.reduce_sum((t_pos-pred_pos) ** 2, axis=-1, keepdims=True) + 1e-4), 0.1) + ploss = tf.debugging.check_numerics(ploss, "ploss loss") + return ploss + + def calc_classification_loss(self, orig_t_pid, pred_id, t_is_unique=None, hasunique=None): + """ + Truth PID is not yet one-hot encoded + Encoding: + 0: Photon + 1: Neutral Hadron + 2: Charged + 3: Ambiguous + """ + if self.classification_loss_weight <= 0: + return tf.reduce_mean(pred_id,axis=1, keepdims=True) + + charged_conditions = [ + tf.abs(orig_t_pid) == 11, # Electrons + tf.abs(orig_t_pid) == 13, # Muons + tf.abs(orig_t_pid) == 211, # Charged Pions + tf.abs(orig_t_pid) == 321, # Charged Kaons + tf.abs(orig_t_pid) == 2212, # Protons + tf.abs(orig_t_pid) == 3112, # Sigma- + tf.abs(orig_t_pid) == 3222, # Sigma+ + tf.abs(orig_t_pid) == 3312, # Xi- + ] + neutral_hadronic_conditions = [ + tf.abs(orig_t_pid) == 130, # Klong + tf.abs(orig_t_pid) == 310, # Kshort + tf.abs(orig_t_pid) == 311, # K0 + tf.abs(orig_t_pid) == 2112, # Neutrons + tf.abs(orig_t_pid) == 3122, # Lambda + tf.abs(orig_t_pid) == 3322, # Xi0 + ] + charged_true = tf.reduce_any(charged_conditions, axis=0) + neutral_hadronic_true = tf.reduce_any(neutral_hadronic_conditions, axis=0) + + truth_pid_tmp = tf.zeros_like(orig_t_pid) - 1 # Initialize with -1 + truth_pid_tmp = tf.where(tf.abs(orig_t_pid) == 22, 0, truth_pid_tmp) # Photons + truth_pid_tmp = tf.where(neutral_hadronic_true, 1, truth_pid_tmp) # Neutral Had. + truth_pid_tmp = tf.where(charged_true, 2, truth_pid_tmp) # Charged + truth_pid_tmp = tf.where(truth_pid_tmp == -1, 3, truth_pid_tmp) # Catch rest + truth_pid_tmp = tf.cast(truth_pid_tmp, tf.int32) + + truth_pid_onehot = tf.one_hot(truth_pid_tmp, depth=4) + truth_pid_onehot = tf.reshape(truth_pid_onehot, (-1, 4)) + + pred_id = tf.clip_by_value(pred_id, 1e-9, 1. - 1e-9) + classloss = tf.keras.losses.categorical_crossentropy(truth_pid_onehot, pred_id) + + classloss = tf.debugging.check_numerics(classloss, "classloss") + + # Weight the loss by the class weights to account for class imbalance + class_weights = tf.constant([1/2.5, 1/0.4, 1/3.75, 1.0], dtype=tf.float32) #Inverse of frequency of classes in COCOA dataset + class_weights = tf.reshape(class_weights, (1, 4)) + weighted_classloss = classloss * tf.reduce_sum(truth_pid_onehot * class_weights, axis=-1) + + weighted_classloss = tf.debugging.check_numerics(weighted_classloss, "weighted_classloss") + + return weighted_classloss[..., tf.newaxis] + class LLExtendedObjectCondensation5(LLExtendedObjectCondensation): ''' diff --git a/modules/OCHits2Showers.py b/modules/OCHits2Showers.py index 1f657413..366d7c73 100644 --- a/modules/OCHits2Showers.py +++ b/modules/OCHits2Showers.py @@ -732,8 +732,9 @@ def process_endcap2(hits2showers_layer, energy_gather_layer, features_dict, """ if not 'no_noise_sel' in predictions_dict.keys(): N_pred = len(predictions_dict['pred_beta']) - predictions_dict['no_noise_sel'] = np.arange(N_pred).reshape((N_pred,1)).astype(int) - is_track = np.abs(features_dict['recHitZ']) == 315 + predictions_dict['no_noise_sel'] = np.arange(N_pred).reshape((N_pred,1)).astype(int) + is_track = np.abs(features_dict['recHitID']) + # is_track = np.abs(features_dict['recHitZ']) == 315 print("Make showers") if not hdbscan: @@ -760,12 +761,12 @@ def process_endcap2(hits2showers_layer, energy_gather_layer, features_dict, if not isinstance(alpha_idx, np.ndarray): alpha_idx = alpha_idx.numpy() alpha_idx = np.reshape(alpha_idx, newshape=(-1,)) - print("Made showers") + # print("Made showers") processed_pred_dict = dict() processed_pred_dict['pred_sid'] = pred_sid - print("Get energy") + # print("Get energy") energy_data = energy_gather_layer( pred_sid, predictions_dict['pred_energy_corr_factor'], @@ -777,7 +778,7 @@ def process_endcap2(hits2showers_layer, energy_gather_layer, features_dict, alpha_idx_tracks = alpha_idx_tracks, t_is_minbias = is_minbias, ) - print("Got energy") + # print("Got energy") try: diff --git a/modules/ShowersMatcher2.py b/modules/ShowersMatcher2.py index aa6cc10c..b2ee7a68 100644 --- a/modules/ShowersMatcher2.py +++ b/modules/ShowersMatcher2.py @@ -150,6 +150,35 @@ def _assign_additional_attr(self): # attr['assigned_to_noise_frac'] = assigned_to_noise_frac # print("A2N", assigned_to_noise_frac) + def _assign_hasTrack(self): + id = self.features_dict['recHitID'] + truth_sid = self.truth_dict['truthHitAssignementIdx'][:, 0].astype(np.int32) * 1 + pred_sid = self.pred_sid * 1 + + for n, attr in self.graph.nodes(data=True): + if attr['type'] == ShowersMatcher._NODE_TYPE_TRUTH_SHOWER: + filt = truth_sid == n + else: + filt = pred_sid == n + + attr['hasTrack'] = np.any(id[filt]) + + #JUST FOR TESTING + def _assign_tidx(self): + truth_sid = self.truth_dict['truthHitAssignementIdx'][:, 0].astype(np.int32) * 1 + pred_sid = self.pred_sid * 1 + + for n, attr in self.graph.nodes(data=True): + if attr['type'] == ShowersMatcher._NODE_TYPE_TRUTH_SHOWER: + filt = truth_sid == n + tidx = self.truth_dict['truthHitAssignementIdx'][filt].flatten() + attr['tidx'] = np.argmax(np.bincount(tidx)) + else: + filt = pred_sid == n + tidx = self.predictions_dict['t_idx'][filt].flatten() + beta = self.predictions_dict['pred_beta'][filt].flatten() + hist, _ = np.histogram(tidx, bins=range(-2, np.max(tidx)+1), weights=beta) + attr['tidx'] = np.argmax(hist)-1 def _build_data_graph(self): """ @@ -201,7 +230,7 @@ def _build_data_graph(self): for k in keys: if k in skip: continue - elif k == 'pred_id': + elif k == 'pred_id' or k == 'pred_pos': # print("SHAPE pred_id: ", self.predictions_dict[k].shape) # print("UNIQUE pred_id: ", np.unique(self.predictions_dict[k])) # node_attributes[k] = np.argmax( @@ -245,13 +274,63 @@ def _cost_matrix_angle_based(self, truth_shower_sid, pred_shower_sid): y['energy']) return C + + def _cost_matrix_Cocoa(self, truth_shower_sid, pred_shower_sid): + + n = max(len(truth_shower_sid), len(pred_shower_sid)) + C = np.zeros((n, n)) + + for a, j in enumerate(pred_shower_sid): + x = self.graph.nodes(data=True)[j] + for b, i in enumerate(truth_shower_sid): + y = self.graph.nodes(data=True)[i] + if (x['hasTrack'] != y['hasTrack']): + C[a, b] =0 + else: + pred_phi = np.arctan2(x['pred_pos'][1], x['pred_pos'][0]) + truth_phi = np.arctan2(y['truthHitAssignedY'], y['truthHitAssignedX']) + deltaPhi = np.abs(pred_phi - truth_phi) + pred_eta = x['pred_pos'][2] + truth_eta = y['truthHitAssignedZ'] + deltaEta = np.abs(pred_eta - truth_eta) + + deltaRsq = deltaPhi**2 + deltaEta**2 + + if y['hasTrack']: + reldeltaptsq = 0 + else: + pt_truth = y['truthHitAssignedEnergies']/np.cosh(truth_eta) + pt_pred = x['pred_energy']/np.cosh(pred_eta) + reldeltaptsq = (np.abs(pt_pred - pt_truth)/ pt_truth)**2 + + d= np.sqrt(reldeltaptsq + 5*deltaRsq) + + C[a, b] = 1/(d+1e-3) + return C + + #JUST FOR TESTING + def _cost_matrix_tidx(self, truth_shower_sid, pred_shower_sid): + + n = max(len(truth_shower_sid), len(pred_shower_sid)) + C = np.zeros((n, n)) + + for a, j in enumerate(pred_shower_sid): + x = self.graph.nodes(data=True)[j] + for b, i in enumerate(truth_shower_sid): + y = self.graph.nodes(data=True)[i] + if (x['tidx'] != y['tidx']): + C[a, b] =0 + else: + C[a, b] = 1 + return C def _cost_matrix_intersection_based(self, truth_shower_sid, pred_shower_sid, return_raw=False): pred_shower_energy = np.array([self.graph.nodes[x]['pred_energy'] for x in pred_shower_sid]) truth_shower_energy = np.array([self.graph.nodes[x]['truthHitAssignedEnergies'] for x in truth_shower_sid]) weight = self.features_dict['recHitEnergy'][:, 0] - is_track = np.abs(self.features_dict['recHitZ'][:,0]) == 315 # TODO: This works for now, but is very hacky + is_track = np.bool8(np.abs(self.features_dict['recHitID'][:,0])) + #is_track = np.abs(self.features_dict['recHitZ'][:,0]) == 315 # TODO: This works for now, but is very hacky weight[is_track] = 0 @@ -329,11 +408,93 @@ def _match_single_pass(self): self.calculated_graph = matched_full_graph + #based on _match_single_pass + def _match_cocoa(self): + truth_shower_sid = [ + x[0] + for x in self.graph.nodes(data=True) + if x[1]['type']==ShowersMatcher._NODE_TYPE_TRUTH_SHOWER + ] + pred_shower_sid = [ + x[0] + for x in self.graph.nodes(data=True) + if x[1]['type']==ShowersMatcher._NODE_TYPE_PRED_SHOWER + ] + + if not len(truth_shower_sid) > 0: + print("No truth showers") + return + if not len(pred_shower_sid) > 0: + print("No predicted showers") + return + + self._assign_hasTrack() + + C = self._cost_matrix_Cocoa(truth_shower_sid, pred_shower_sid) + + row_id, col_id = linear_sum_assignment(C, maximize=True) + + self.truth_sid = self.truth_dict['truthHitAssignementIdx'][:, 0] + matched_full_graph = nx.Graph() + matched_full_graph.add_nodes_from(self.graph.nodes(data=True)) + + for p, t in zip(row_id, col_id): + if C[p, t] > 0: + matched_full_graph.add_edge( + truth_shower_sid[t], + pred_shower_sid[p], + attached_in_pass=0) + + self.calculated_graph = matched_full_graph + + #JUST FOR TESTING + def _match_tidx(self): + truth_shower_sid = [ + x[0] + for x in self.graph.nodes(data=True) + if x[1]['type']==ShowersMatcher._NODE_TYPE_TRUTH_SHOWER + ] + pred_shower_sid = [ + x[0] + for x in self.graph.nodes(data=True) + if x[1]['type']==ShowersMatcher._NODE_TYPE_PRED_SHOWER + ] + + if not len(truth_shower_sid) > 0: + print("No truth showers") + return + if not len(pred_shower_sid) > 0: + print("No predicted showers") + return + + self._assign_tidx() + + C = self._cost_matrix_tidx(truth_shower_sid, pred_shower_sid) + + row_id, col_id = linear_sum_assignment(C, maximize=True) + + self.truth_sid = self.truth_dict['truthHitAssignementIdx'][:, 0] + matched_full_graph = nx.Graph() + matched_full_graph.add_nodes_from(self.graph.nodes(data=True)) + + for p, t in zip(row_id, col_id): + if C[p, t] > 0: + matched_full_graph.add_edge( + truth_shower_sid[t], + pred_shower_sid[p], + attached_in_pass=0) + + self.calculated_graph = matched_full_graph def process(self, extra=False): self._build_data_graph() if self.match_mode == 'iou_max': self._match_single_pass() + elif self.match_mode == 'cocoa': + self._match_cocoa() + elif self.match_mode == 'tidx': + #JUST FOR TESTING + self._match_tidx() else: raise NotImplementedError('Match mode not found') if extra: diff --git a/modules/compiled/bin_by_coordinates_kernel.cc b/modules/compiled/bin_by_coordinates_kernel.cc index f0c0cb1e..80c13348 100644 --- a/modules/compiled/bin_by_coordinates_kernel.cc +++ b/modules/compiled/bin_by_coordinates_kernel.cc @@ -23,12 +23,12 @@ struct BinByCoordinatesNbinsHelperOpFunctor { //just because a int n_nbins, int nrs){ int n=1; - printf("n bins:"); + //printf("n bins:"); for(int i=0;i { const int *d_truthidx, const int *d_unique_idx, + const int *rs, int * out_idx, - float * m_not, + int * m_not, const int n_vert, const int n_unique, const int n_max_per_unique, + const int n_max_in_rs, + + const int n_rs, bool calc_m_not) { //main axis: n_unique == K_obj @@ -53,7 +57,7 @@ struct MIndicesOpFunctor { for (int k = 0; k < n_unique; k++) { - // + int uqidx = d_unique_idx[k]; int puqcounter=0; @@ -66,15 +70,33 @@ struct MIndicesOpFunctor { for(int prem = puqcounter; prem < n_max_per_unique; prem++){ out_idx[I2D(k, prem, n_max_per_unique)] = -1; } + //m_not + if(calc_m_not){ + //find row for uqidx + int rowForUqidx = 0; + while(rowForUqidx+1= rs[rowForUqidx+1]){ + rowForUqidx++; + } - // - if(calc_m_not ){ - //m_not + + int mnot_index = 0; for(int i_v = 0; i_v < n_vert; i_v++ ){ - if(uqidx>=0 && d_truthidx[i_v] == uqidx) - m_not [I2D(k, i_v, n_vert)] = 0.; - else - m_not [I2D(k, i_v, n_vert)] = 1.; + //find row for i_v + int rowFori_v= 0; + while(rowFori_v+1= rs[rowFori_v+1]){ + rowFori_v++; + } + + //compare rows and index + if (rowFori_v == rowForUqidx && (uqidx < 0 || d_truthidx[i_v] != uqidx)){ + m_not [I2D(k, mnot_index, n_max_in_rs)] = i_v; + mnot_index++; + } + } + //fill rest with -1 + while(mnot_index < n_max_in_rs){ + m_not [I2D(k, mnot_index, n_max_in_rs)] = -1; + mnot_index++; } } @@ -103,17 +125,27 @@ class MIndicesOp : public OpKernel { const Tensor &t_tidx = context->input(0); const Tensor &t_uqtixs = context->input(1); const Tensor &t_nmax_puq = context->input(2);//needs to be evaluated in kernel + const Tensor &t_rs = context->input(3); + const Tensor &t_max_in_rs = context->input(4); const int n_vert = t_tidx.dim_size(0); const int n_unique = t_uqtixs.dim_size(0); - int n_max_per_unique=0; + const int n_rs = t_rs.dim_size(0); + int n_max_per_unique=0; MIndicesMaxUqOpFunctor()( context->eigen_device(), t_nmax_puq.flat().data(), &n_max_per_unique ); + int n_max_in_rs = 0; + MIndicesMaxUqOpFunctor()( + context->eigen_device(), + t_max_in_rs.flat().data(), + &n_max_in_rs + ); + Tensor *out_idx = NULL; OP_REQUIRES_OK(context, context->allocate_output(0, {n_unique, n_max_per_unique}, &out_idx)); @@ -121,8 +153,7 @@ class MIndicesOp : public OpKernel { Tensor *m_not = NULL; TensorShape m_not_shape={1, 1}; if(calc_m_not_) - m_not_shape={n_unique, n_vert}; - + m_not_shape={n_unique, n_max_in_rs}; OP_REQUIRES_OK(context, context->allocate_output(1, m_not_shape, &m_not)); @@ -130,13 +161,16 @@ class MIndicesOp : public OpKernel { context->eigen_device(), t_tidx.flat().data(), t_uqtixs.flat().data(), + t_rs.flat().data(), out_idx->flat().data(), - m_not->flat().data(), + m_not->flat().data(), n_vert, n_unique, n_max_per_unique, + n_rs, + n_max_in_rs, calc_m_not_ ); diff --git a/modules/compiled/oc_helper_m_indices_kernel.cu.cc b/modules/compiled/oc_helper_m_indices_kernel.cu.cc index eb303f6e..f91536ca 100644 --- a/modules/compiled/oc_helper_m_indices_kernel.cu.cc +++ b/modules/compiled/oc_helper_m_indices_kernel.cu.cc @@ -22,13 +22,17 @@ __global__ static void calc( const int *d_truthidx, const int *d_unique_idx, + const int *rs, int * out_idx, - float * m_not, + int * m_not, const int n_vert, const int n_unique, const int n_max_per_unique, + const int n_rs, + const int n_max_in_rs, + bool calc_m_not){ const int k = blockIdx.x * blockDim.x + threadIdx.x; @@ -49,15 +53,33 @@ static void calc( for(int prem = puqcounter; prem < n_max_per_unique; prem++){ out_idx[I2D(k, prem, n_max_per_unique)] = -1; } + //m_not + if(calc_m_not){ + //find row for uqidx + int rowForUqidx = 0; + while(rowForUqidx+1= rs[rowForUqidx+1]){ + rowForUqidx++; + } - // - if(calc_m_not ){ - //m_not + + int mnot_index = 0; for(int i_v = 0; i_v < n_vert; i_v++ ){ - if(uqidx>=0 && d_truthidx[i_v] == uqidx) - m_not [I2D(k, i_v, n_vert)] = 0.; - else - m_not [I2D(k, i_v, n_vert)] = 1.; + //find row for i_v + int rowFori_v= 0; + while(rowFori_v+1= rs[rowFori_v+1]){ + rowFori_v++; + } + + //compare rows and index + if (rowFori_v == rowForUqidx && (uqidx < 0 || d_truthidx[i_v] != uqidx)){ + m_not [I2D(k, mnot_index, n_max_in_rs)] = i_v; + mnot_index++; + } + } + //fill rest with -1 + while(mnot_index < n_max_in_rs){ + m_not [I2D(k, mnot_index, n_max_in_rs)] = -1; + mnot_index++; } } @@ -86,13 +108,16 @@ struct MIndicesOpFunctor { const int *d_truthidx, const int *d_unique_idx, + const int *rs, int * out_idx, - float * m_not, + int * m_not, const int n_vert, const int n_unique, const int n_max_per_unique, + const int n_rs, + const int n_max_in_rs, bool calc_m_not @@ -102,6 +127,7 @@ struct MIndicesOpFunctor { calc<<>>( d_truthidx, d_unique_idx, + rs, out_idx, m_not, @@ -109,6 +135,8 @@ struct MIndicesOpFunctor { n_vert, n_unique, n_max_per_unique, + n_rs, + n_max_in_rs, calc_m_not ); diff --git a/modules/compiled/oc_helper_m_indices_kernel.h b/modules/compiled/oc_helper_m_indices_kernel.h index dda3275d..e9893630 100644 --- a/modules/compiled/oc_helper_m_indices_kernel.h +++ b/modules/compiled/oc_helper_m_indices_kernel.h @@ -21,13 +21,16 @@ struct MIndicesOpFunctor { const int *d_truthidx, const int *d_unique_idx, + const int *rs, int * out_idx, - float * m_not, + int * m_not, const int n_vert, const int n_unique, const int n_max_per_unique, + const int n_rs, + const int n_max_in_rs, bool calc_m_not diff --git a/modules/compiled/oc_helper_m_indices_module.cc b/modules/compiled/oc_helper_m_indices_module.cc index f20af945..5218f49d 100644 --- a/modules/compiled/oc_helper_m_indices_module.cc +++ b/modules/compiled/oc_helper_m_indices_module.cc @@ -18,7 +18,9 @@ REGISTER_OP("MIndices") .Input("asso_idxs: int32") .Input("unique_idxs: int32") .Input("nmax_per_unique: int32") + .Input("rs: int32") + .Input("nmax_in_rs: int32") .Output("sel_idxs: int32") - .Output("m_not: float32"); + .Output("m_not: int32"); diff --git a/modules/datastructures/TrainData_Cocoa.py b/modules/datastructures/TrainData_Cocoa.py new file mode 100644 index 00000000..9f6e6e30 --- /dev/null +++ b/modules/datastructures/TrainData_Cocoa.py @@ -0,0 +1,208 @@ +""" +Class for the conversion of the COCOA data to DJC format +start conversion with 'convertDJCFromSource -i /path/to/input -o /path/to/output -c TrainData_Cocoa --gpu' +""" +from datastructures.TrainData_NanoML import TrainData_NanoML +import uproot +import pandas as pd +import awkward as ak +import numpy as np +from DeepJetCore import SimpleArray + + +class TrainData_Cocoa(TrainData_NanoML): + + def convertFromSourceFile(self, filename, weighterobjects, istraining): + #open File + with uproot.open(filename) as file: + events = file[file.keys()[0]] + data = events.arrays(library='ak')#[0:10]#JUST FOR TESTING REMOVE SLICING LATER!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + + #convert data to training data + trainData,rs = self.converttotrainingdfvec(data) + + #make simpleArray features + features = trainData[['recHitEnergy', 'recHitEta', 'isTrack','recHitTheta', 'recHitR', 'recHitX', 'recHitY', 'recHitZ', 'recHitTime', 'recHitHitR']] + farr = SimpleArray(np.float32(features.to_numpy()),rs,name="recHitFeatures") + + #make simpleArray truth + t={} + #First the float32 fields + for field in ['t_energy','t_spectator', 't_time','t_rec_energy']: + t[field] = SimpleArray(np.float32(trainData[field].to_numpy().reshape(-1, 1)),rs,name=field) + #Then the int32 fields + for field in ['t_idx', 't_pid', 't_fully_contained', 't_is_unique']: + t[field] = SimpleArray(np.int32(trainData[field].to_numpy().reshape(-1, 1)),rs,name=field) + #Then the pos field + t['t_pos'] = SimpleArray(np.float32(np.array(trainData['t_pos'].tolist())),rs,name='t_pos') + + #return result + return [farr, + t['t_idx'], t['t_energy'], t['t_pos'], t['t_time'], + t['t_pid'], t['t_spectator'], t['t_fully_contained'], + t['t_rec_energy'], t['t_is_unique'] ],[], [] + + def convertevent(self, event): + #convert cells to df + df_cell = pd.DataFrame() + df_cell['recHitEnergy'] = ak.to_dataframe(event['cell_e'], how='outer')/1000 + df_cell['recHitEta'] = ak.to_dataframe(event['cell_eta'], how='outer') + df_cell['isTrack'] = 0 + df_cell['recHitTheta'] = 0 #will be calculated later + df_cell['recHitR'] = 0 #will be calculated later + df_cell['recHitX']= ak.to_dataframe(event['cell_x'], how='outer') + df_cell['recHitY']= ak.to_dataframe(event['cell_y'], how='outer') + df_cell['recHitZ']= ak.to_dataframe(event['cell_z'], how='outer') + df_cell['recHitTime'] = 0 + df_cell['recHitHitR'] = 0 + + df_cell['t_idx'] = ak.to_dataframe(event['cell_parent_idx'], how='outer') + + #Calculate Distance from 0,0,0 + df_cell['recHitR']= np.sqrt(df_cell['recHitX']**2+df_cell['recHitY']**2+df_cell['recHitZ']**2) + df_cell['recHitTheta'] = np.arctan2(df_cell['recHitR'], df_cell['recHitZ']) + + + #convert particle information to df + df_particle= pd.DataFrame() + df_particle['particle_idx'] = np.arange(len(event['particle_e'])) + df_particle['t_energy'] = ak.to_dataframe(event['particle_e'], how='outer')/1000 + df_particle['t_pid'] = ak.to_dataframe(event['particle_pdgid'], how='outer') + df_particle['t_rec_energy'] = ak.to_dataframe(event['particle_dep_energy'], how='outer')/1000 + df_particle['particle_eta'] = ak.to_dataframe(event['particle_eta_lay0'], how='outer') + df_particle['particle_phi'] = ak.to_dataframe(event['particle_phi_lay0'], how='outer') + + #convert track information to df + df_track = pd.DataFrame() + df_track['t_idx'] = ak.to_dataframe(event['track_parent_idx'], how='outer') + df_track['recHitX'] = ak.to_dataframe(event['track_x_layer_0'], how='outer') + df_track['recHitY'] = ak.to_dataframe(event['track_y_layer_0'], how='outer') + df_track['recHitZ'] = ak.to_dataframe(event['track_z_layer_0'], how='outer') + df_track['recHitR']= np.sqrt(df_track['recHitX']**2+df_track['recHitY']**2+df_track['recHitZ']**2) + + df_track['recHitTheta'] = np.arctan2(df_track['recHitR'], df_track['recHitZ']) + df_track['recHitEta'] = ak.to_dataframe(event['track_eta_layer_0'], how='outer') + + df_track['recHitTime'] = 0 + df_track['recHitHitR'] = 0 + + #join track information to particle information + df_track = df_track.join(df_particle, on='t_idx') + + #set is Track to -1 or 1 depending on charge of particle + df_track['isTrack'] = pdgid_to_charge(df_track['t_pid']) + #raise Error if any isTrack is zero (should not happen) + if 0 in df_track['isTrack'].values: + raise ValueError("Error in isTrack. Unknown PDGID for particle with track.") + + #Smear true energy for track + df_track['recHitEnergy'] = df_track['t_energy'].apply(lambda x: np.random.normal(x, 0.01 * x)) + + #join cell information to truth information + df_cell = df_cell.join(df_particle, on='t_idx') + + #concat cell and track information + df_training = pd.concat([df_cell, df_track], ignore_index=True) + + #set NaN values (particle information for hits made from background noise to -1) + df_training.fillna(-1, inplace=True) + + noisemask = df_training['t_idx']==-1 + + #Set other default values + df_training['t_time'] = 0 + df_training['t_spectator'] = 0 + df_training['t_fully_contained'] = 1 + + #set t_is_unique + df_training['t_is_unique'] = 0 + first_occurrence_mask = (df_training['t_idx'] != -1) & (df_training.groupby('t_idx').cumcount() == 0) + df_training.loc[first_occurrence_mask, 't_is_unique'] = 1 + + #encode pos as cos(phi), sin(phi), eta + df_training['t_pos'] = df_training.apply(lambda row: [np.cos(row['particle_phi']), np.sin(row['particle_phi']), row['particle_eta']], axis=1) + + #set t_pos to [cos(phi), sin(phi), eta] based on RecHitX,Y,Z if t_idx is -1 + df_training.loc[noisemask, 't_pos'] = df_training[noisemask].apply(lambda row: [np.cos(np.arctan2(row['recHitY'], row['recHitX'])), np.sin(np.arctan2(row['recHitY'], row['recHitX'])), row['recHitEta']], axis=1) + + #replace -1 in t_energy and t_rec_energy with the sum of all noise hits + noiseEnergy = np.sum(df_training[noisemask]['recHitEnergy']) + df_training.loc[noisemask, 't_energy'] = noiseEnergy + df_training.loc[noisemask, 't_rec_energy'] = noiseEnergy + + #set is_energymax to one if the hit has the maximum energy for each particle (t_idx) + # df_training['is_energymax'] = 0 + # for i in np.unique(df_training['t_idx']): + # mask = np.logical_and(df_training['t_idx'] == i, df_training['t_idx'] != -1) + # max_energy_idx = df_training.loc[mask, 'recHitEnergy'].idxmax() + # df_training.loc[mask, 'is_energymax'][max_energy_idx] = 1 + + # #JUST FOR TESTING FIXME + # df_training['isTrack']=df_training['is_energymax'] + + return df_training + + def converttotrainingdfvec(self, data): + #For some reason, there is an outlier event with a track at -1e12, which is not physical + print("Number of events before removing outlier: ", len(data)) + data = data[ak.all(data.track_x_layer_0 > -2000, axis=1)] + print("Number of events after removing outlier: ", len(data)) + + #Convert events one by one + traindata = np.array([self.convertevent(data[eventnumber]) for eventnumber in np.arange(len(data))], dtype=object) + + #Remove broken events (set to False for Testdata) + if True: + mask = np.ones(len(traindata), dtype=bool) + #Remove all events with an energy lower than 15GeV + # E_cutoff = 15000 + # E_sum = ak.sum(data['particle_e'], axis=1) + # mask = ak.to_numpy(E_sum >= E_cutoff) + + #Remove all events with a particle that has multiple clusters + for i in range(len(data)): + for j in range(len(data.particle_pdgid[i])): + particlemask = data.cell_parent_idx[i] == j + mean_phi = ak.mean(data.cell_phi[i][particlemask]) + mean_eta = ak.mean(data.cell_eta[i][particlemask]) + phi_mask = abs(data.cell_phi[i][particlemask] - mean_phi) < 0.5 + eta_mask = abs(data.cell_eta[i][particlemask] - mean_eta) < 0.5 + energy_outside = ak.sum(data.cell_e[i][particlemask][~phi_mask | ~eta_mask]) + energy_total = ak.sum(data.cell_e[i][particlemask]) + mask[i] = mask[i] and energy_outside < 0.1 * energy_total + + + print("Number of events before cuts: ", len(traindata)) + traindata = traindata[mask] + print("Number of events after cuts: ", len(traindata)) + + #find the row splits + rs = np.cumsum([len(df) for df in traindata]) + rs = np.insert(rs, 0, 0) + + #concat all events + traindata = pd.concat(traindata) + return traindata, rs + +#HelpFunction to convert PDGIDs to the charge of the particle +def pdgid_to_charge(pdgid): + charge_dict = { + 11: -1, # electron + -11: 1, # positron + 13: -1, # muon + -13: 1, # antimuon + 211: 1, # pi+ + -211: -1, # pi- + 321: 1, # K+ + -321: -1, # K- + 2212: 1, # proton + -2212: -1, # antiproton + 3112: -1, # Sigma- + -3112: 1, # antilambda_c+ + 3222: 1, # Sigma+ + -3222: -1, # antilambda_c- + 3312: -1, # Xi- + -3312: 1, # antixi+ + } + # Default to 0, check if valid later + return np.vectorize(charge_dict.get)(pdgid, 0) diff --git a/modules/extra_plots.py b/modules/extra_plots.py index 2acc337e..588e48f8 100644 --- a/modules/extra_plots.py +++ b/modules/extra_plots.py @@ -236,6 +236,7 @@ def energy_resolution(df, bins=None, binwidth=10., addfit=False): # ax2.hist(bins[1:], weights=entries, histtype='stepfilled', color='grey', alpha=0.2, zorder=1) ax2.hist(bins[1:], weights=entries, color='grey', alpha=0.2, zorder=1, label='Counts') yticks = ax2.get_yticks() + ax2.set_yticks(yticks) ax2.set_yticklabels(yticks, fontsize=20) ax2.set_title("Respone and Counts", fontsize=40) ax2.set_ylabel("Counts", fontsize=40) @@ -254,6 +255,7 @@ def energy_resolution(df, bins=None, binwidth=10., addfit=False): ax3.set_xticks(xticks) ax3.set_xticklabels(xticks, fontsize=20) yticks = ax3.get_yticks() + ax3.set_yticks(yticks) ax3.set_yticklabels(np.round(yticks, 3), fontsize=20) ax3.grid(alpha=0.5, linestyle='--') # fit resolution function to sigmaOverE @@ -572,6 +574,7 @@ def noise_performance(noise_df): return fig MAP_DICT = { + -1: 5, # Noise 0: 5, # Whatever 13: 0, # Muon -13: 0, diff --git a/modules/model_blocks.py b/modules/model_blocks.py index a5d8d4ea..5b4e7eec 100644 --- a/modules/model_blocks.py +++ b/modules/model_blocks.py @@ -426,7 +426,8 @@ def create_outputs(x, n_ccoords=3, trainable=True, is_track=None, set_track_betas_to_one=False, - predict_spectator_weights=False): + predict_spectator_weights=False, + pred_e_factor = 0.01): ''' returns * pred_beta Dense(1) @@ -477,7 +478,7 @@ def create_outputs(x, n_ccoords=3, name = name_prefix+'_energy', kernel_initializer='zeros', activation=energy_act, - trainable=trainable)(ScalarMultiply(0.01)(x)) + trainable=trainable)(ScalarMultiply(pred_e_factor)(x)) if energy_factor: pred_energy = Add(name= name_prefix+'_one_plus_energy')([OnesLike()(pred_energy),pred_energy]) diff --git a/modules/object_condensation.py b/modules/object_condensation.py index 59dbfc73..a3f2a1ae 100644 --- a/modules/object_condensation.py +++ b/modules/object_condensation.py @@ -1,11 +1,10 @@ # -*- coding: utf-8 -*- import tensorflow as tf -from oc_helper_ops import CreateMidx, SelectWithDefault +from oc_helper_ops import CreateMidx, SelectWithDefault, SelectWithDefaultMnot from binned_select_knn_op import BinnedSelectKnn - def huber(x, d): losssq = x**2 absx = tf.abs(x) @@ -118,8 +117,8 @@ def _weighted_alpha_k(self, v_k_m, w_k_m, fraction_wrt_alpha : float): return ak * (1. - fraction_wrt_alpha) + fraction_wrt_alpha * wvk - def create_Ms(self, truth_idx): - self.Msel, self.Mnot, _ = CreateMidx(truth_idx, calc_m_not=True) + def create_Ms(self, truth_idx, rs): + self.Msel, self.Mnot, _ = CreateMidx(truth_idx, rs, calc_m_not=True) def set_input(self, @@ -128,6 +127,7 @@ def set_input(self, d, pll, truth_idx, + rs, object_weight, is_spectator_weight, calc_Ms=True, @@ -153,7 +153,7 @@ def set_input(self, self.q_v = tf.where(truth_idx < 0, noise_qmin, self.q_v) # set noise qmin if calc_Ms: - self.create_Ms(truth_idx) + self.create_Ms(truth_idx, rs=rs) if self.Msel is None: self.valid=False return @@ -226,13 +226,13 @@ def calc_dsq_rep(self): dsq = tf.expand_dims(self.x_k, axis=1) - tf.expand_dims(self.x_v, axis=0) #K x V x C dsq = tf.reduce_sum(dsq**2, axis=-1, keepdims=True) #K x V x 1 return dsq - + #@tf.function def V_rep_k(self): K = tf.reduce_sum(tf.ones_like(self.q_k)) - N_notk = tf.reduce_sum(self.Mnot, axis=1) + #future remark: if this gets too large, one could use a kNN/radiusNN here dsq = self.calc_dsq_rep() @@ -240,11 +240,18 @@ def V_rep_k(self): # nogradbeta = tf.stop_gradient(self.beta_k_m) #weight. tf.reduce_sum( tf.exp(-dsq) * d_v_e, , axis=1) / tf.reduce_sum( tf.exp(-dsq) ) sigma = self.weighted_d_k_m(dsq) #create gradients for all, but prefer k vertex - dsq = tf.math.divide_no_nan(dsq, sigma**2 + 1e-4) #K x V x 1 + dsq = tf.math.divide_no_nan(dsq, sigma**2 + 1e-4) + - V_rep = self.rep_func(dsq) * self.Mnot * tf.expand_dims(self.q_v,axis=0) #K x V x 1 + V_rep = self.rep_func(dsq) * tf.expand_dims(self.q_v,axis=0) + V_rep = SelectWithDefaultMnot(self.Mnot, V_rep) #K x V x 1 + + + dsq_mnot = SelectWithDefaultMnot(self.Mnot, dsq) #K x V x 1 if self.rep_range > 0: - N_notk = tf.reduce_sum(tf.where(dsq < self.rep_range**2 , 1., tf.zeros_like(dsq)) * self.Mnot, axis=1) + N_notk = tf.reduce_sum(tf.where(dsq_mnot < self.rep_range**2 , 1., tf.zeros_like(dsq_mnot)), axis=1) + else: + N_notk = tf.reduce_sum(tf.ones_like(dsq_mnot), axis=1) V_rep = self.q_k * tf.reduce_sum(V_rep, axis=1) #K x 1 #if self.global_weight: @@ -261,6 +268,7 @@ def V_rep_k(self): def Pll_k(self): tanhsqbeta = self.beta_v**2 #softer here + #tanhsqbeta = tf.math.atanh(self.beta_v*(1-1e-3))**2 tanhsqbeta = tf.debugging.check_numerics(tanhsqbeta, "OC: pw b**2") pw = tanhsqbeta * tf.clip_by_value((1.-tf.clip_by_value(self.isn_v+self.sw_v,0.,1.)),0.,1.) + 1e-6 @@ -394,15 +402,17 @@ def _calc_contamination(self, energies, rel_metrics_radius, energies_k): x_k_alpha = tf.gather_nd(self.x_k_m,self.alpha_k, batch_dims=1) dsq = tf.expand_dims(x_k_alpha, axis=1) - tf.expand_dims(self.x_v, axis=0) #K x V x C dsq = tf.reduce_sum(dsq**2, axis=-1) #K x V - in_radius = rel_metrics_radius > tf.sqrt(dsq) / self.d_k# K x V + in_radius = rel_metrics_radius > tf.sqrt(dsq) / self.d_k # K x V + in_radius_mnot = SelectWithDefaultMnot(self.Mnot, in_radius) #K x V energies_k_v = tf.expand_dims(energies, axis=0) # K x V x 1 energies_ir_all_k_v = tf.where(in_radius[...,tf.newaxis], energies_k_v , 0.) - energies_ir_not_k_v = tf.where(in_radius[...,tf.newaxis], self.Mnot * energies_k_v , 0.) + + energies_ir_not_k_v = tf.where(in_radius_mnot[...,tf.newaxis], energies_k_v , 0.) rel_cont_k = tf.reduce_sum(energies_ir_not_k_v, axis=1)/tf.reduce_sum(energies_ir_all_k_v, axis=1) - - beta_cont_k = tf.reduce_sum(tf.where(in_radius[...,tf.newaxis], self.Mnot * self.beta_v[tf.newaxis,...], 0.), axis=1) # K x 1? + + beta_cont_k = tf.reduce_sum(tf.where(in_radius_mnot[...,tf.newaxis], self.beta_v[tf.newaxis,...], 0.), axis=1) # K x 1? beta_cont_k /= tf.reduce_sum(tf.where(in_radius[...,tf.newaxis], self.beta_v[tf.newaxis,...], 0.), axis=1) # the same for deposited energy > 20 GeV @@ -445,6 +455,29 @@ def __init__(self, **kwargs): super(Hinge_OC_per_sample, self).__init__(**kwargs) self.condensation_damping = 0.0 # Don't stop any gradients #self.rep_range = 2. +class Hinge_OC_per_sample_atanhweighing(Hinge_OC_per_sample): + ''' + Same as Hinge_OC_per_sample but with atanh(beta)^2 weighting for the payloads + ''' + def __init__(self, **kwargs): + super(Hinge_OC_per_sample_atanhweighing, self).__init__(**kwargs) + #@tf.function + def Pll_k(self): + #tanhsqbeta = self.beta_v**2 #softer here + tanhsqbeta = tf.math.atanh(self.beta_v*(1-1e-3))**2 + tanhsqbeta = tf.debugging.check_numerics(tanhsqbeta, "OC: pw b**2") + pw = tanhsqbeta * tf.clip_by_value((1.-tf.clip_by_value(self.isn_v+self.sw_v,0.,1.)),0.,1.) + 1e-6 + + pw = tf.debugging.check_numerics(pw, "OC: pw") + + pll_k_m = SelectWithDefault(self.Msel, self.pll_v, 0.) #K x V_perobj x P + pw_k_m = SelectWithDefault(self.Msel, pw, 0.) #K x V-obj x P + pw_k_sum = tf.reduce_sum(pw_k_m, axis=1) + pw_k_sum = tf.where(pw_k_sum <= 0., 1e-2, pw_k_sum) + + pll_k = tf.math.divide_no_nan(tf.reduce_sum(pll_k_m * pw_k_m, axis=1), + pw_k_sum )#K x P + return pll_k @@ -677,9 +710,9 @@ def set_input(self, beta, ): #replace beta with per-object normalised value - self.create_Ms(truth_idx) - beta_max = tf.reduce_max(self.Mnot * tf.expand_dims(beta,axis=0),axis=1, keepdims=True) #K x 1 x 1 - beta_max *= self.Mnot #K x V x 1 + self.create_Ms(truth_idx, rs=rs) + beta_max = tf.reduce_max(SelectWithDefaultMnot(self.Mnot, tf.expand_dims(beta,axis=0),0),axis=1, keepdims=True) #K x 1 x 1 + beta_max = SelectWithDefault(self.Msel, beta_max, 0.) # this is the same reduced in K given the others are zero with exception of noise (filtered later) beta_max = tf.reduce_max(beta_max, axis=0) # V x 1 @@ -710,10 +743,8 @@ def __call__(self, beta, truth_idx, object_weight, is_spectator_weight, - rs, energies = None): #rs last - tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen = 6*[tf.constant(0., tf.float32)] #batch loop @@ -721,29 +752,23 @@ def __call__(self, beta, return tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen batch_size = rs.shape[0] - 1 mdict = {} - - for b in tf.range(batch_size): - - self.loss_impl.set_input( - beta[rs[b]:rs[b + 1]], - x[rs[b]:rs[b + 1]], - d[rs[b]:rs[b + 1]], - pll[rs[b]:rs[b + 1]], - truth_idx[rs[b]:rs[b + 1]], - object_weight[rs[b]:rs[b + 1]], - is_spectator_weight[rs[b]:rs[b + 1]] + self.loss_impl.set_input( + beta, + x, + d, + pll, + truth_idx, + rs, + object_weight, + is_spectator_weight ) - tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen = self.loss_impl.add_to_terms( - tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen - ) - - #just last row split is good enough for metric - if energies is not None and b == batch_size-1: - mdict = self.loss_impl.calc_metrics(energies[rs[b]:rs[b + 1]]) - - bs = tf.cast(batch_size, dtype='float32') + 1e-3 - out = [a/bs for a in [tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen]] + tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen = self.loss_impl.add_to_terms( + tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen + ) + + mdict = self.loss_impl.calc_metrics(energies) + out = [tot_V_att, tot_V_rep, tot_Noise_pen, tot_B_pen, tot_pll,tot_too_much_B_pen] return out, mdict diff --git a/modules/oc_helper_ops.py b/modules/oc_helper_ops.py index d9685614..e24f39ca 100644 --- a/modules/oc_helper_ops.py +++ b/modules/oc_helper_ops.py @@ -28,7 +28,7 @@ _op = tf.load_op_library('oc_helper_m_indices.so') -def CreateMidx(truth_idxs, calc_m_not=False): +def CreateMidx(truth_idxs, rs, calc_m_not=False): ''' /* * Takes as helper input @@ -50,8 +50,14 @@ def CreateMidx(truth_idxs, calc_m_not=False): ''' - #only consider non-noise - c_truth_idxs = truth_idxs[truth_idxs>=0] + # Repeat each rs according to the differences to offset the indices + differences = tf.concat([tf.math.abs(rs[1:] - rs[:-1]), [0]], axis=0) + repeated_rs = tf.reshape(tf.repeat(rs, differences),(-1,1)) + offset_truth_idxs = truth_idxs + repeated_rs + + offset_truth_idxs = tf.where(truth_idxs<0, -1, offset_truth_idxs) + + c_truth_idxs= offset_truth_idxs[offset_truth_idxs>=0] unique_idxs, _, cperunique = tf.unique_with_counts(c_truth_idxs) nmax_per_unique = tf.reduce_max(cperunique) @@ -59,15 +65,20 @@ def CreateMidx(truth_idxs, calc_m_not=False): #eager if nmax_per_unique.numpy() < 1: return None, None, None - + + nmax_in_rs = tf.reduce_max(rs[1:]-rs[:-1]) + sel_dxs, m_not = _op.MIndices( calc_m_not=calc_m_not, - asso_idxs = truth_idxs, + asso_idxs = offset_truth_idxs, unique_idxs = unique_idxs, - nmax_per_unique = nmax_per_unique + nmax_per_unique = nmax_per_unique, + rs = rs, + nmax_in_rs = nmax_in_rs ) - return sel_dxs, tf.expand_dims(m_not,axis=2), tf.expand_dims(cperunique,axis=1) #just some conventions + + return sel_dxs, m_not, tf.expand_dims(cperunique,axis=1) #just some conventions @ops.RegisterGradient("CreateMidx") def _CreateMidxGrad(op, sel_dxs, m_not): @@ -90,6 +101,43 @@ def SelectWithDefault(indices, tensor, default=0, no_check=False): tf.assert_equal(tf.shape(tensor)[1], tf.shape(out)[2])]): return out +def SelectWithDefaultMnot(Mnot, tensor): + # Check if tensor has a third dimension, if not, expand it + no_extra_dim = (len(tf.shape(tensor)) == 2) + if no_extra_dim: + tensor = tf.expand_dims(tensor, axis=-1) + + K, V, _ = tensor.shape + + # Create a base mask of shape (K, V) initialized to False + mask = tf.zeros((K, V), dtype=tf.bool) + + # Flatten Mnot and create row indices + row_indices = tf.repeat(tf.range(K), repeats=Mnot.shape[1]) + col_indices = tf.reshape(Mnot, [-1]) + + # Filter out invalid indices (-1) + valid_mask = col_indices >= 0 + row_indices = tf.boolean_mask(row_indices, valid_mask) + col_indices = tf.boolean_mask(col_indices, valid_mask) + + # Stack indices together + indices = tf.stack([row_indices, col_indices], axis=1) + + # Scatter update to create the mask + updates = tf.ones(indices.shape[0], dtype=tf.bool) + mask = tf.tensor_scatter_nd_update(mask, indices, updates) + + # Expand the mask to match the last dimension of the tensor + mask_expanded = tf.expand_dims(mask, -1) + + # Apply the mask to set the elements to 0 + tensor = tf.where(mask_expanded, tensor, tf.zeros_like(tensor)) + + if no_extra_dim: + tensor = tf.squeeze(tensor, axis=-1) + + return tensor def per_rs_segids_to_unique(pred_sid, rs, return_nseg=False, strict_check=True): diff --git a/scripts/analyse_cocoa_predictions.py b/scripts/analyse_cocoa_predictions.py new file mode 100755 index 00000000..5aabbf24 --- /dev/null +++ b/scripts/analyse_cocoa_predictions.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +""" +Analysis script adapted from analyse_hgcal_predictions to run of the predictions of the model for the COCOA dataset. +Call with python analyse_cocoa_predictions.py -h for help. +""" + +import os +import argparse +import pickle +import gzip +import mgzip +import pandas as pd +import numpy as np +import tensorflow as tf +import matplotlib.pyplot as plt +import fastjet as fj +from scipy.optimize import linear_sum_assignment + +from OCHits2Showers import OCHits2ShowersLayer +from ShowersMatcher2 import ShowersMatcher +from plot_cocoa import plot_everything + + +def analyse(preddir, + outputpath, + beta_threshold, + distance_threshold, + iou_threshold, + matching_mode, + analysisoutpath, + nfiles, + nevents, + local_distance_scaling, + de_e_cut, + angle_cut, + extra=False, + shower0=False, + datasetname='Quark Jet' + ): + + hits2showers = OCHits2ShowersLayer( + beta_threshold, + distance_threshold, + local_distance_scaling) + showers_matcher = ShowersMatcher(matching_mode, iou_threshold, de_e_cut, angle_cut, shower0) + + + files_to_be_tested = [ + os.path.join(preddir, x) + for x in os.listdir(preddir) + if (x.endswith('.bin.gz') and x.startswith('pred'))] + if nfiles!=-1: + files_to_be_tested = files_to_be_tested[0:min(nfiles, len(files_to_be_tested))] + + showers_dataframe = pd.DataFrame() + jets_dataframe = pd.DataFrame() + features = [] + truth = [] + prediction = [] + alpha_ids = [] + matched = [] + event_id = 0 + + ############################################################################################### + ### Loop over all events ###################################################################### + ############################################################################################### + + for i, file in enumerate(files_to_be_tested): + print(f"Analysing file {i+1}/{len(files_to_be_tested)}") + with mgzip.open(file, 'rb') as analysis_file: + file_data = pickle.load(analysis_file) + for j, endcap_data in enumerate(file_data): + if (nevents != -1) and (j > nevents): + continue + # print(f"Analysing endcap {j+1}/{len(file_data)}") + features_dict, truth_dict, predictions_dict = endcap_data + features.append(features_dict) + truth.append(truth_dict) + + print(f"Analyzing event {event_id}") + + #Create showers + pred_sid, _, alpha_idx, _, _ = hits2showers( + predictions_dict['pred_ccoords'], + predictions_dict['pred_beta'], + predictions_dict['pred_dist']) + + predictions_dict['pred_sid'] = pred_sid + prediction.append(predictions_dict) + + #Match showers + if not isinstance(alpha_idx, np.ndarray): + alpha_idx = alpha_idx.numpy() + alpha_idx = np.reshape(alpha_idx, newshape=(-1,)) + + alpha_ids.append(alpha_idx) + + showers_matcher.set_inputs( + features_dict=features_dict, + truth_dict=truth_dict, + predictions_dict=predictions_dict, + pred_alpha_idx=alpha_idx + ) + showers_matcher.process(extra=extra) + dataframe = showers_matcher.get_result_as_dataframe() + matched_truth_sid, matched_pred_sid = showers_matcher.get_matched_hit_sids() + matched.append((matched_truth_sid, matched_pred_sid)) + + dataframe['event_id'] = event_id + dataframe['pred_id_value'] = dataframe['pred_id'].apply(lambda x: np.argmax(x)) + showers_dataframe = pd.concat((showers_dataframe, dataframe)) + + ############################################################################################### + ### Jet Clustering ############################################################################ + ############################################################################################### + + has_truth = np.isnan(dataframe['truthHitAssignedEnergies']) == False + has_pred = np.isnan(dataframe['pred_energy']) == False + + E_truth = np.asarray(dataframe[has_truth]['truthHitAssignedEnergies'], dtype=np.double) + eta_truth = dataframe[has_truth]['truthHitAssignedZ'] + phi_truth = np.arctan2(dataframe[has_truth]['truthHitAssignedY'], dataframe[has_truth]['truthHitAssignedX']) + pt_truth = E_truth / np.cosh(eta_truth) + px_truth = np.asarray(pt_truth * np.cos(phi_truth), dtype=np.double) + py_truth = np.asarray(pt_truth * np.sin(phi_truth), dtype=np.double) + pz_truth = np.asarray(pt_truth * np.sinh(eta_truth), dtype=np.double) + + particles_truth = np.array([fj.PseudoJet(px_truth[i], py_truth[i], pz_truth[i], E_truth[i]) for i in range(len(dataframe[has_truth]['truthHitAssignedX']))]) + + E_pred = np.asarray(dataframe[has_pred]['pred_energy'], dtype=np.double) + eta_pred = dataframe[has_pred]['pred_pos'].apply(lambda x: x[2]) + phi_pred = dataframe[has_pred]['pred_pos'].apply(lambda x: np.arctan2(x[1], x[0])) + pt_pred = E_pred / np.cosh(eta_pred) + px_pred = np.asarray(pt_pred * np.cos(phi_pred), dtype=np.double) + py_pred = np.asarray(pt_pred * np.sin(phi_pred), dtype=np.double) + pz_pred = np.asarray(pt_pred * np.sinh(eta_pred), dtype=np.double) + particles_pred = np.array([fj.PseudoJet(px_pred[i], py_pred[i], pz_pred[i], E_pred[i]) for i in range(len(dataframe[has_pred]['pred_pos']))]) + + # Choose jet clustering algorithm and parameters + R = 0.4 # Jet radius parameter + jet_def = fj.JetDefinition(fj.antikt_algorithm, R) + + # Cluster jets + cluster_sequence_truth = fj.ClusterSequence(particles_truth.tolist(), jet_def) + cluster_sequence_pred = fj.ClusterSequence(particles_pred.tolist(), jet_def) + jets_truth = [jet for jet in cluster_sequence_truth.inclusive_jets() if len(jet.constituents()) >= 2] + jets_pred = [jet for jet in cluster_sequence_pred.inclusive_jets() if len(jet.constituents()) >= 2] + + jet_truth_data = { + 'true_pt': [jet.pt() for jet in jets_truth], + 'true_eta': [jet.eta() for jet in jets_truth], + 'true_phi': [jet.phi() for jet in jets_truth], + 'true_mass': [jet.m() for jet in jets_truth], + 'true_n_constituents': [len(jet.constituents()) for jet in jets_truth] + } + jet_pred_dta = { + 'pred_pt': [jet.pt() for jet in jets_pred], + 'pred_eta': [jet.eta() for jet in jets_pred], + 'pred_phi': [jet.phi() for jet in jets_pred], + 'pred_mass': [jet.m() for jet in jets_pred], + 'pred_n_constituents': [len(jet.constituents()) for jet in jets_pred] + } + jet_truth_df = pd.DataFrame(jet_truth_data) + jet_pred_df = pd.DataFrame(jet_pred_dta) + + jet_truth_df['event_id'] = event_id + jet_pred_df['event_id'] = event_id + + # Match jets + n = max(len(jet_truth_df), len(jet_pred_df)) + C = np.zeros((n, n)) + + for a in range(len(jet_truth_df)): + for b in range(len(jet_pred_df)): + delta_eta = jet_truth_df['true_eta'].iloc[a] - jet_pred_df['pred_eta'].iloc[b] + delta_phi = jet_truth_df['true_phi'].iloc[a] - jet_pred_df['pred_phi'].iloc[b] + C[a, b] = 1/np.sqrt(delta_eta**2 + delta_phi**2) + row_id, col_id = linear_sum_assignment(C, maximize=True) + + matched_jets = [] + for a, b in zip(row_id, col_id): + if C[a, b] > 0: + matched_jets.append({ + 'matched': True, + 'true_pt': jet_truth_df['true_pt'].iloc[a], + 'true_eta': jet_truth_df['true_eta'].iloc[a], + 'true_phi': jet_truth_df['true_phi'].iloc[a], + 'true_mass': jet_truth_df['true_mass'].iloc[a], + 'true_n_constituents': jet_truth_df['true_n_constituents'].iloc[a], + 'pred_pt': jet_pred_df['pred_pt'].iloc[b], + 'pred_eta': jet_pred_df['pred_eta'].iloc[b], + 'pred_phi': jet_pred_df['pred_phi'].iloc[b], + 'pred_mass': jet_pred_df['pred_mass'].iloc[b], + 'pred_n_constituents': jet_pred_df['pred_n_constituents'].iloc[b] + }) + else: + if C[a, :].max() == 0 and a 0: + + scalar_variables = { + 'beta_threshold': str(beta_threshold), + 'distance_threshold': str(distance_threshold), + 'iou_threshold': str(iou_threshold), + 'matching_mode': str(matching_mode), + 'de_e_cut': str(de_e_cut), + 'angle_cut': str(angle_cut), + } + + analysis_data = { + 'showers_dataframe' : showers_dataframe, + 'jets_dataframe' : jets_dataframe, + 'scalar_variables' : scalar_variables, + 'alpha_ids' : alpha_ids, + 'matched': matched, + 'features': features, + 'truth': truth, + 'prediction': prediction + } + + with gzip.open(analysisoutpath, 'wb') as output_file: + print("Writing dataframes to pickled file",analysisoutpath) + pickle.dump(analysis_data, output_file) + + ############################################################################################### + ### Make Plots ################################################################################ + ############################################################################################### + + if(len(outputpath) > 0): + print('Starting to plot', outputpath) + plot_everything(showers_dataframe, prediction, truth, features, jets_dataframe, outputpath, datasetname) + print("DONE") + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + 'Analyse predictions from object condensation and plot relevant results') + parser.add_argument('preddir', + help='Directory with .bin.gz files or a txt file with full paths of the \ + bin-gz files from the prediction.') + parser.add_argument('-p', + help="Output directory for the final analysis pdf file (otherwise, it won't be produced)", + default='') + parser.add_argument('-b', help='Beta threshold (default 0.1)', default='0.1') + parser.add_argument('-d', help='Distance threshold (default 0.5)', default='0.5') + parser.add_argument('-i', help='IOU threshold (default 0.1)', default='0.1') + parser.add_argument('-m', help='Matching mode', default='cocoa') + parser.add_argument('--analysisoutpath', + help='Will dump analysis data to a file to remake plots without re-running everything.', + default='') + parser.add_argument('--nfiles', + help='Maximum number of files. -1 for everything in the preddir', + default=-1) + parser.add_argument('--nevents', help='Maximum number of events (per file)', default=-1) + parser.add_argument('--no_local_distance_scaling', help='With local distance scaling', + action='store_true') + parser.add_argument('--de_e_cut', help='dE/E threshold to allow match.', default=-1) + parser.add_argument('--angle_cut', help='Angle cut for angle based matching', default=-1) + parser.add_argument('--extra', + help="Calculate more information for showers", + action='store_true') + parser.add_argument('--shower0', + help="Only match with truth shower ID=0", + action='store_true') + parser.add_argument('--gluon', action='store_true', help='Writing Gluon dataset') + + args = parser.parse_args() + + if args.gluon: + datasetname='Gluon Jet' + else: + datasetname='Quark Jet' + + analyse(preddir=args.preddir, + outputpath=args.p, + beta_threshold=float(args.b), + distance_threshold=float(args.d), + iou_threshold=float(args.i), + matching_mode=args.m, + analysisoutpath=args.analysisoutpath, + nfiles=int(args.nfiles), + local_distance_scaling=not args.no_local_distance_scaling, + de_e_cut=float(args.de_e_cut), + angle_cut=float(args.angle_cut), + nevents=int(args.nevents), + extra=args.extra, + shower0=args.shower0, + datasetname=datasetname, + ) + diff --git a/scripts/eval_cocoa.py b/scripts/eval_cocoa.py new file mode 100644 index 00000000..75b648e9 --- /dev/null +++ b/scripts/eval_cocoa.py @@ -0,0 +1,87 @@ +""" +Helper script to run prediction and analysis for a cocoa model on both test sets. +""" + +from DeepJetCore.wandb_interface import wandb_wrapper as wandb +wandb.active = False + +from argparse import ArgumentParser +from hgcal_predictor import HGCalPredictor +import os +from analyse_cocoa_predictions import analyse + +if __name__ == '__main__': + parser = ArgumentParser( + 'Helper script to run prediction and analysis for a cocoa model on both test sets.') + parser.add_argument('modellocation', + help='location of the KERAS_model.h5 file') + parser.add_argument('outputdir', + help="Output directory", + default='') + parser.add_argument('-testfiles', + help="Location of a folder with a quarkDC._n.djcdc and gluonDC._n.djcdc test file", + default='/work/friemer/hgcalml/testdata/') + parser.add_argument("--max_files", help="Limit number of files", default=-1) + parser.add_argument("--max_steps", help="Limit number of steps per file", default=-1) + #Arguments for analyse + parser.add_argument('-b', help='Beta threshold (default 0.1)', default='0.1') + parser.add_argument('-d', help='Distance threshold (default 0.5)', default='0.5') + parser.add_argument('-i', help='IOU threshold (default 0.1)', default='0.1') + parser.add_argument('-m', help='Matching mode', default='cocoa') + parser.add_argument('--nfiles', + help='Maximum number of files to analyse. -1 for everything in the preddir', + default=-1) + parser.add_argument('--nevents', help='Maximum number of events (per file) to analyse', default=-1) + + args = parser.parse_args() + + +HGCalPredictor( + os.path.join(args.testfiles, 'quarkDC._n.djcdc'), + os.path.join(args.testfiles, 'quarkDC._n.djcdc'), + os.path.join(args.outputdir, 'predictionQuark'), + inputdir=None, + unbuffered=False, + max_files=int(args.max_files), + toydata=False, + ).predict(model_path=args.modellocation, max_steps=int(args.max_steps)) + +HGCalPredictor( + os.path.join(args.testfiles, 'gluonDC._n.djcdc'), + os.path.join(args.testfiles, 'gluonDC._n.djcdc'), + os.path.join(args.outputdir, 'predictionGluon'), + inputdir=None, + unbuffered=False, + max_files=int(args.max_files), + toydata=False, + ).predict(model_path=args.modellocation, max_steps=int(args.max_steps)) + +analyse(preddir=os.path.join(args.outputdir, 'predictionQuark'), + outputpath=os.path.join(args.outputdir, 'plotsQuark'), + beta_threshold=float(args.b), + distance_threshold=float(args.d), + iou_threshold=float(args.i), + matching_mode=args.m, + analysisoutpath=os.path.join(args.outputdir, 'analysisQuark'), + nfiles=int(args.nfiles), + local_distance_scaling=True, + de_e_cut=-1.0, + angle_cut=-1.0, + nevents=int(args.nevents), + datasetname='Quark Jet', + ) + +analyse(preddir=os.path.join(args.outputdir, 'predictionGluon'), + outputpath=os.path.join(args.outputdir, 'plotsGluon'), + beta_threshold=float(args.b), + distance_threshold=float(args.d), + iou_threshold=float(args.i), + matching_mode=args.m, + analysisoutpath=os.path.join(args.outputdir, 'analysisGluon'), + nfiles=int(args.nfiles), + local_distance_scaling=True, + de_e_cut=-1.0, + angle_cut=-1.0, + nevents=int(args.nevents), + datasetname='Gluon Jet', + ) \ No newline at end of file diff --git a/scripts/plotWindow.py b/scripts/plotWindow.py new file mode 100644 index 00000000..384ef4bb --- /dev/null +++ b/scripts/plotWindow.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 + +from DeepJetCore import DataCollection +import os +from argparse import ArgumentParser +import numpy as np +import matplotlib.pyplot as plt +import math +from multiprocessing import Process +import random +from datastructures import TrainData_NanoML, TrainData_NanoMLPF, TrainData_Cocoa +import plotly.express as px +import pandas as pd +import tqdm +from djcdata.dataPipeline import TrainDataGenerator + + +parser = ArgumentParser('') +parser.add_argument('inputFile') +parser.add_argument('outputDir') +parser.add_argument('--hipsearch',action='store_true') +parser.add_argument('--plots',action='store_true') +parser.add_argument('--pf',action='store_true') +args = parser.parse_args() + +outdir = args.outputDir+'/' +### rewrite! +os.system('mkdir -p '+outdir) +td_class = TrainData_Cocoa +#td_class = TrainData_NanoML +if args.pf: + td_class = TrainData_NanoMLPF + +#read a file +def invokeGen(infile): + if infile[-6:] == '.djcdc': + dc = DataCollection(infile) + td = dc.dataclass() + tdclass = dc.dataclass + dc.setBatchSize(1) + gen = dc.invokeGenerator() + elif infile[-6:] == '.djctd': + td = td_class() + tdclass = td_class + td.readFromFile(infile) + gen = TrainDataGenerator() + gen.setBatchSize(1) + gen.setBuffer(td) + elif infile[-5:] == '.root': + print('reading from root file, converting...') + td = td_class() + tdclass = td_class + td.readFromSourceFile(infile,{},True) + td.writeToFile(infile+'.djctd') + td.readFromFile(infile+'.djctd') + print('conversion done') + gen = TrainDataGenerator() + gen.setBatchSize(1) + gen.setBuffer(td) + + gen.setSkipTooLargeBatches(False) + nevents = gen.getNBatches() + #gen.cast_to = tdclass + gen.cast_to = None + return gen.feedTrainData,nevents,td + +gen,nevents,td = invokeGen(args.inputFile) + +def shuffle_truth_colors(df, qualifier="truthHitAssignementIdx"): + ta = df[qualifier] + unta = np.unique(ta) + unta = unta[unta>-0.1] + np.random.seed(42) + np.random.shuffle(unta) + out = ta.copy() + dfo = df.copy() + for i in range(len(unta)): + out[ta ==unta[i]]=i + dfo[qualifier] = out + return dfo + +def toDataFrame(thegen, thetd): + data = next(thegen())#this is a dict, row splits can be ignored, this is per event + return TrainData_NanoML.createPandasDataFrame(data, 0) + + +def compressShowerFeatures(df): + dfout = df.drop_duplicates(subset = ["truthHitAssignementIdx"]) + return dfout[dfout["truthHitAssignementIdx"]>=0] + +def quickplotshower(df,out): + fig = px.scatter_3d(df, x="recHitX", y="recHitZ", z="recHitY", + color="hitratio", size="recHitLogEnergy", + symbol = "marker", + hover_data=['rel_std','totruthHitAssignedEnergies_ratio','marker','hitratio','nhits','corratio'], + template='plotly_dark', + color_continuous_scale=px.colors.sequential.Rainbow) + fig.update_traces(marker=dict(line=dict(width=0))) + fig.write_html(out) + +def hipsearch(df3d, i, outdir, makeplots=False): + truthHitAssignementIdx = df3d['truthHitAssignementIdx'] + utidx = np.unique(truthHitAssignementIdx) + counter=0 + Er_dep=[] + Er_corr_dep=[] + E =[] + for t in utidx: + if t < 0: + continue + seldf = df3d[df3d['truthHitAssignementIdx']==t] + + depsum = np.ones_like(seldf['recHitEnergy'])*np.sum(seldf['recHitEnergy']) + nhits = float(len(seldf['recHitEnergy'])) + seldf['energy_ratio'] = seldf['recHitEnergy']/seldf['truthHitAssignedEnergies'] + + seldf['totruthHitAssignedEnergies_ratio'] = depsum/seldf['truthHitAssignedEnergies'] + E.append(np.mean(seldf['truthHitAssignedEnergies'])) + + Er_dep.append(np.mean(seldf['totruthHitAssignedEnergies_ratio'])) + + hitratio = seldf['recHitEnergy']/depsum + seldf['nhits'] = nhits + hitratio *= nhits #1. on average for uniform etc. + seldf['hitratio'] = hitratio + + m = np.mean(seldf['hitratio']) + s = np.std(seldf['hitratio']-m) + + seldf['rel_std']= (hitratio-m)/s + seldf['marker'] = np.array(seldf['rel_std'] > 5.,dtype='int32') + ewithout = np.sum((1.-seldf['marker'])*seldf['recHitEnergy']) + seldf['corratio'] = ewithout/seldf['truthHitAssignedEnergies'] + + Er_corr_dep.append(np.mean(seldf['corratio'])) + if makeplots and np.all(depsum < seldf['truthHitAssignedEnergies']*1.1): + quickplotshower(seldf,outdir+str(i)+'_'+str(counter)+'.html') + counter+=1 + + return Er_dep, Er_corr_dep , E + + + +hitdf = pd.DataFrame() +showerdf = pd.DataFrame() +eventdf = pd.DataFrame() + +print(nevents,'events') +#3D plots +Er_dep, Er_corr_dep, E = [], [], [] +for i in tqdm.tqdm(range(nevents)): + + df = toDataFrame(gen,td) + + #print(df.columns) + + dfshowers = compressShowerFeatures(df) + showerhits = df[df["truthHitAssignementIdx"]>=0] + #depvstruthenergy.append(np.sum(showerhits['recHitEnergy'])/(np.sum(dfshowers['truthHitAssignedEnergies'])+1.)) + + from globals import pu + #df["recHitLogEnergy"]*= (1. - (1.-1e-2)*(df["truthHitAssignementIdx"]>=pu.t_idx_offset)) + df3d = shuffle_truth_colors(df) + df3d['orig_truthHitAssignementIdx']=df['truthHitAssignementIdx'] + df3d['t_inv_spec'] = np.where(df3d['truthHitAssignementIdx']<0, + 0.,np.ones_like(df3d['truthHitSpectatorFlag'])) * 1./(df3d['truthHitSpectatorFlag']+1e-1) + #plot the first 20 as 3D plots + if i < 20 and args.plots: + #makes a copy + + hover_data=['recHitEnergy', + 'recHitHitR', + 'truthHitAssignedEnergies', + 'truthHitAssignedT', + 'truthHitAssignedX', + 'truthHitAssignedY', + 'truthHitAssignedZ', + 'truthHitAssignementIdx', + 'orig_truthHitAssignementIdx', + 'truthHitAssignedPIDs', + 'truthHitSpectatorFlag'] + + print('N hits', len(df3d)) + + fig = px.scatter_3d(df3d, x="recHitX", y="recHitZ", z="recHitY", + color="truthHitAssignementIdx", size="recHitLogEnergy", + symbol = "recHitID", + hover_data=hover_data, + template='plotly_dark', + color_continuous_scale=px.colors.sequential.Rainbow) + fig.update_traces(marker=dict(line=dict(width=0))) + ccfile = outdir + str(i) + "_event.html" + fig.write_html(ccfile) + + continue + + fig = px.scatter_3d(df3d, x="recHitX", y="recHitZ", z="recHitY", + color="truthHitAssignementIdx", size="recHitHitR", + symbol = "recHitID", + hover_data=hover_data, + template='plotly_dark', + color_continuous_scale=px.colors.sequential.Rainbow) + fig.update_traces(marker=dict(line=dict(width=0))) + ccfile = outdir + str(i) + "_event_hitsize.html" + fig.write_html(ccfile) + + fig = px.scatter_3d(df3d, x="recHitX", y="recHitZ", z="recHitY", + color="truthHitAssignementIdx", size="t_inv_spec", + hover_data=hover_data, + symbol = "recHitID", + template='plotly_dark', + color_continuous_scale=px.colors.sequential.Rainbow) + fig.update_traces(marker=dict(line=dict(width=0))) + ccfile = outdir + str(i) + "_spect.html" + fig.write_html(ccfile) + + if args.hipsearch: + iEr_dep, iEr_corr_dep, iE = hipsearch(df3d, i, outdir, args.plots) + Er_dep+=iEr_dep + Er_corr_dep+=iEr_corr_dep + E+=iE + +plt.hist(Er_dep,bins=31,label='uncorr',alpha=0.5) +plt.hist(Er_corr_dep,bins=31,label='corr',alpha=0.5) +plt.legend() +plt.savefig(outdir +'hipcorr.pdf') + + +dfout = pd.DataFrame(zip(E,Er_dep,Er_corr_dep), columns = ['E','Er','Er_corr']) +dfout.to_pickle("df.pkl") \ No newline at end of file diff --git a/scripts/plot_cocoa.py b/scripts/plot_cocoa.py new file mode 100644 index 00000000..e61e4e9f --- /dev/null +++ b/scripts/plot_cocoa.py @@ -0,0 +1,941 @@ +""" +Script to create evaluation plots for COCOA models +either call plot_everything function from different script like analyse_cocoa_predictions.py +or run this script directly with the following command if analysisfile is already created: +python plot_cocoa.py /path/to/analysisfile /path/to/outputdir +""" + +import gzip +import pickle +import matplotlib.pyplot as plt +import matplotlib +import numpy as np +import pandas as pd +import tensorflow as tf +import sys +import os +import plotly.express as px +import argparse +from scipy.stats import norm +from scipy.optimize import curve_fit + + +def calc_efficiencies(df, bins, mask=None): + """ + Helper function to calculate the efficiencies, fake rates and classification accuracy for the photons and neutral hadrons + """ + efficiencies = [] + efficiencies_err = [] + fake_rate = [] + fake_rate_err = [] + corr_class_prob = [] + corr_class_prob_err = [] + + mask_predicted = np.isnan(df['pred_energy']) == False + mask_truth = np.isnan(df['truthHitAssignedEnergies']) == False + + if mask == None: + mask_PID_truth = np.ones(len(df), dtype=bool) + mask_PID_pred = mask_PID_truth + else: + if(mask == 0): + mask_PID_truth = df['truthHitAssignedPIDs'].isin([22]) + elif(mask == 1): + mask_PID_truth = df['truthHitAssignedPIDs'].isin([130,310,311,2112,-2112,3122,-3122,3322,-3322]) + else: + raise ValueError("mask must be 0 or 1") + mask_PID_pred = df['pred_id_value'].isin([mask]) + + mask_PID_matched = np.logical_and(mask_PID_truth, mask_PID_pred) + + + for i in range(len(bins) - 1): + mask_bintruth = np.logical_and( + df['truthHitAssignedEnergies'] >= bins[i], + df['truthHitAssignedEnergies'] < bins[i + 1]) + mask_binpred = np.logical_and( + df['pred_energy'] >= bins[i], + df['pred_energy'] < bins[i + 1]) + + matched = np.logical_and( + mask_predicted, + mask_bintruth) + faked = np.logical_and( + mask_binpred, + np.logical_not(mask_truth)) + + eff = np.sum(matched[mask_PID_truth]) / np.sum(mask_bintruth[mask_PID_truth]) + eff_err = np.sqrt(eff * (1 - eff) / np.sum(matched[mask_PID_truth])) + efficiencies.append(eff) + efficiencies_err.append(eff_err) + + fake = np.sum(faked[mask_PID_pred]) / np.sum(mask_binpred[mask_PID_pred]) + fake_err = np.sqrt(fake * (1 - fake) / np.sum(faked[mask_PID_pred])) + fake_rate.append(fake) + fake_rate_err.append(fake_err) + + cc_prob = np.sum(matched[mask_PID_matched]) / np.sum(matched[mask_PID_truth]) + cc_prob_err = np.sqrt(cc_prob * (1 - cc_prob) / np.sum(matched[mask_PID_truth])) + corr_class_prob.append(cc_prob) + corr_class_prob_err.append(cc_prob_err) + + return np.array(efficiencies), np.array(efficiencies_err), np.array(fake_rate), np.array(fake_rate_err), np.array(corr_class_prob), np.array(corr_class_prob_err) + +def plot_efficencies(df, bins): + """ + Function to plot the efficiencies, fake rates and classification accuracy for the photons and neutral hadrons similar to the HGPFlow paper + """ + # Calculate the efficiencies and fake rates for the photons + yeff_photon, yerr_eff_photon, yfake_photon, yerr_fake_photon, ycorr_photon, yerr_corr_photon = calc_efficiencies(df, bins, 0) + yeff_photon, yerr_eff_photon, yfake_photon, yerr_fake_photon, ycorr_photon, yerr_corr_photon = \ + yeff_photon*100, yerr_eff_photon*100, yfake_photon*100, yerr_fake_photon*100, ycorr_photon*100, yerr_corr_photon*100 + + # Calculate the efficiencies and fake rates for the neutral hadrons + #130: "K0L", 310: "K0S", 311: "K0", 2112: "neutron",-2112: "antineutron",3122: "Lambda",-3122: "antilambda"3322: "Xi0",-3322: "antixi0" + yeff_nh, yerr_eff_nh, yfake_nh, yerr_fake_nh, ycorr_nh, yerr_corr_nh = calc_efficiencies(df, bins, 1) + yeff_nh, yerr_eff_nh, yfake_nh, yerr_fake_nh, ycorr_nh, yerr_corr_nh = \ + yeff_nh*100, yerr_eff_nh*100, yfake_nh*100, yerr_fake_nh*100, ycorr_nh*100, yerr_corr_nh*100 + + # Calculate the bin positions and widths + binwidth = bins[1:] - bins[:-1] + x_pos = bins[:-1] + binwidth / 2 + x_err = binwidth / 2 + + # Create the plots + fig, ((ax2, ax3), (ax5,ax6), (ax8,ax9)) = plt.subplots(nrows=3, ncols=2, figsize=(12, 10)) + ax2.errorbar(x_pos, yeff_photon, xerr=x_err, yerr=yerr_eff_photon, fmt='o', label='Efficiency') + ax2.set_ylabel('efficiency gamma', fontsize=10) + ax3.errorbar(x_pos, yeff_nh, xerr=x_err, yerr=yerr_eff_nh, fmt='o', label='Efficiency') + ax3.set_ylabel('efficiency n.H.', fontsize=10) + ytickseff = np.round(np.arange(40, 101, 20), 1) + ax2.set_yticks(ytickseff) + ax2.set_yticklabels([f"{y}%" for y in ytickseff], fontsize=10) + ax3.set_yticks(ytickseff) + ax3.set_yticklabels([f"{y}%" for y in ytickseff], fontsize=10) + + ax5.errorbar(x_pos, yfake_photon, xerr=x_err, yerr=yerr_fake_photon, fmt='o', label='Fake rate') + ax5.set_ylabel('fake rate gamma', fontsize=10) + ax6.errorbar(x_pos, yfake_nh, xerr=x_err, yerr=yerr_fake_nh, fmt='o', label='Fake rate') + ax6.set_ylabel('fake rate n.H.', fontsize=10) + yticksfake1 = np.round(np.arange(0, 50, 20), 1) + ax5.set_yticks(yticksfake1) + ax5.set_yticklabels([f"{y}%" for y in yticksfake1], fontsize=10) + yticksfake2 = np.round(np.arange(0, 70, 20), 1) + ax6.set_yticks(yticksfake2) + ax6.set_yticklabels([f"{y}%" for y in yticksfake2], fontsize=10) + + ax8.errorbar(x_pos, ycorr_photon, xerr=x_err, yerr=yerr_corr_photon, fmt='o', label='probability of correct class') + ax8.set_ylabel('p corr gamma', fontsize=10) + ax9.errorbar(x_pos, ycorr_nh, xerr=x_err, yerr=yerr_corr_nh, fmt='o', label='probability of correct class') + ax9.set_ylabel('p corr n.H.', fontsize=10) + yticksp = np.round(np.arange(0, 101, 20), 1) + ax8.set_yticks(yticksp) + ax8.set_yticklabels([f"{y}%" for y in yticksp], fontsize=10) + ax9.set_yticks(yticksp) + ax9.set_yticklabels([f"{y}%" for y in yticksp], fontsize=10) + + for ax in [ax2,ax3, ax5, ax6, ax8, ax9]: + ax.set_xticks(bins) + ax.set_xticklabels(bins, fontsize=10) + ax.set_xlim(bins[0], bins[-1]) + ax.grid(alpha=0.4) + ax.set_xlabel('Energy [GeV]', fontsize=10) + #ax.set_ylim(20, 101) + + return fig +def plot_efficency_and_fakerate(df, bins, datasetname='Quark Jet'): + # Calculate the efficiencies and fake rates for the total + efficiencies, efficiencies_err, fake_rate, fake_rate_err = [], [], [], [] + + mask_predicted = np.isnan(df['pred_energy']) == False + mask_truth = np.isnan(df['truthHitAssignedEnergies']) == False + + + for i in range(len(bins) - 1): + mask_bintruth = np.logical_and( + df['truthHitAssignedEnergies'] >= bins[i], + df['truthHitAssignedEnergies'] < bins[i + 1]) + mask_binpred = np.logical_and( + df['pred_energy'] >= bins[i], + df['pred_energy'] < bins[i + 1]) + + matched = np.logical_and( + mask_predicted, + mask_bintruth) + faked = np.logical_and( + mask_binpred, + np.logical_not(mask_truth)) + + eff = np.sum(matched) / np.sum(mask_bintruth) + eff_err = np.sqrt(eff * (1 - eff) / np.sum(matched)) + efficiencies.append(eff) + efficiencies_err.append(eff_err) + + fake = np.sum(faked) / np.sum(mask_binpred) + fake_err = np.sqrt(fake * (1 - fake) / np.sum(faked)) + fake_rate.append(fake) + fake_rate_err.append(fake_err) + + yeff, yerr_eff, yfake, yerr_fake, = np.array(efficiencies), np.array(efficiencies_err), np.array(fake_rate), np.array(fake_rate_err) + + # Calculate the bin positions and widths + binwidth = bins[1:] - bins[:-1] + x_pos = bins[:-1] + binwidth / 2 + x_err = binwidth / 2 + + fig_eff = plt.figure() + plt.errorbar(x_pos, yeff, xerr=x_err, yerr=yerr_eff, fmt='o') + plt.xlim(bins[0], bins[-1]) + plt.xticks(np.array([1,5,10,20,30,50])) + plt.xlabel('Truth Energy [GeV]') + plt.ylabel('Efficiency') + plt.grid(alpha=0.4) + plt.text(0.05, 0.95, datasetname, horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + plt.close() + + fig_fake = plt.figure() + plt.errorbar(x_pos, yfake, xerr=x_err, yerr=yerr_fake, fmt='o') + plt.xlim(bins[0], bins[-1]) + plt.xticks(np.array([1,5,10,20,30,50])) + plt.xlabel('Predicted Energy [GeV]') + plt.ylim(bottom=0, top=max(0.1, np.max(yfake) + 0.01)) + plt.ylabel('Fake rate') + plt.grid(alpha=0.4) + plt.text(0.05, 0.95, datasetname, horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + plt.close() + + return fig_eff, fig_fake + +def calc_classification_p(df, bins, mask): + """ + Helper function to calculate the classification accuracy for the photons, neutral hadrons and charged particles + """ + corr_class_prob = [] + corr_class_prob_err = [] + + mask_predicted = np.isnan(df['pred_energy']) == False + mask_truth = np.isnan(df['truthHitAssignedEnergies']) == False + + if(mask == 0): + mask_PID_truth = df['truthHitAssignedPIDs'].isin([22]) + elif(mask == 1): + mask_PID_truth = df['truthHitAssignedPIDs'].isin([130,310,311,2112,-2112,3122,-3122,3322,-3322]) + elif(mask == 2): + mask_PID_truth = df['truthHitAssignedPIDs'].isin([11,-11,13,-13,211,-211,321,-321,2212,-2212,3112,-3112,3222,-3222,3312,-3312]) + else: + raise ValueError("mask must be 0, 1 or 2") + mask_PID_pred = df['pred_id_value'].isin([mask]) + mask_PID_matched = np.logical_and(mask_PID_truth, mask_PID_pred) + + for i in range(len(bins) - 1): + mask_bintruth = np.logical_and( + df['truthHitAssignedEnergies'] >= bins[i], + df['truthHitAssignedEnergies'] < bins[i + 1]) + mask_binpred = np.logical_and( + df['pred_energy'] >= bins[i], + df['pred_energy'] < bins[i + 1]) + + matched = np.logical_and( + mask_predicted, + mask_bintruth) + + cc_prob = np.sum(matched[mask_PID_matched]) / np.sum(matched[mask_PID_truth]) + cc_prob_err = np.sqrt(cc_prob * (1 - cc_prob) / np.sum(matched[mask_PID_truth])) + corr_class_prob.append(cc_prob) + corr_class_prob_err.append(cc_prob_err) + + return np.array(corr_class_prob), np.array(corr_class_prob_err) + +def plot_classification_p(df, bins, datasetname='Quark Jet'): + ycorr_photon, yerr_corr_photon = calc_classification_p(df, bins, 0) + ycorr_nh, yerr_corr_nh = calc_classification_p(df, bins, 1) + ycorr_ch, yerr_corr_ch = calc_classification_p(df, bins, 2) + + # Calculate the bin positions and widths + binwidth = bins[1:] - bins[:-1] + x_pos = bins[:-1] + binwidth / 2 + x_err = binwidth / 2 + + # Create the plots + fig_classification = plt.figure() + plt.errorbar(x_pos, ycorr_photon, xerr=x_err, yerr=yerr_corr_photon, fmt='o', label='Photon') + plt.errorbar(x_pos, ycorr_nh, xerr=x_err, yerr=yerr_corr_nh, fmt='o', label='Neutral Hadron') + plt.errorbar(x_pos, ycorr_ch, xerr=x_err, yerr=yerr_corr_ch, fmt='o', label='Charged Particle') + + plt.legend() + plt.grid(alpha=0.4) + plt.xlim(bins[0], bins[-1]) + plt.xticks([1,5,10,20,30,50]) + plt.xlabel('Energy [GeV]') + plt.ylabel('Accuracy of classification') + plt.ylim(0, 1) + plt.text(0.05, 0.95, datasetname, horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + plt.close() + + return fig_classification + +def plot_jet_metrics_based_on_particles(df, datasetname='QuarkJet'): + """ + plots the jet metrics under the assumption, that one event is one jet + """ + + reseta = [] + resphi = [] + relresE = [] + nParicles_pred = [] + nParicles_truth = [] + true_pt = [] + pred_pt = [] + + has_truth = np.isnan(df['truthHitAssignedEnergies']) == False + has_pred = np.isnan(df['pred_energy']) == False + + for i in range(len(df['event_id'].unique())): + event_mask = df['event_id'] == i + truth_mask = np.logical_and(event_mask, has_truth) + pred_mask = np.logical_and(event_mask, has_pred) + + jet_E_truth = np.sum(df[truth_mask]['truthHitAssignedEnergies']) + jet_E_pred = np.sum(df[pred_mask]['pred_energy']) + relresE.append((jet_E_truth-jet_E_pred)/jet_E_truth) + + nParicles_pred.append(np.sum(pred_mask)) + nParicles_truth.append(np.sum(truth_mask)) + + t_eta = np.average(df[truth_mask]['truthHitAssignedZ'], weights=df[truth_mask]['truthHitAssignedEnergies']) + pred_eta = np.average(df[pred_mask]['pred_pos'].apply(lambda x: x[2]), weights=df[pred_mask]['pred_energy']) + reseta.append(t_eta - pred_eta) + + t_phi = np.average(np.arctan2(df[truth_mask]['truthHitAssignedY'], df[truth_mask]['truthHitAssignedX']), weights=df[truth_mask]['truthHitAssignedEnergies']) + pred_phi = np.average(df[pred_mask]['pred_pos'].apply(lambda x: np.arctan2(x[1], x[0])), weights=df[pred_mask]['pred_energy']) + delta_phi = calc_deltaphi(t_phi, pred_phi) + resphi.append(delta_phi) + + true_pt.append(np.sum(df[truth_mask]['truthHitAssignedEnergies'] / np.cosh(df[truth_mask]['truthHitAssignedZ']))) + pred_pt.append(np.sum(df[pred_mask]['pred_energy'] / np.cosh(df[pred_mask]['pred_pos'].apply(lambda x: x[2])))) + + #Correct jet pt + jets_df = pd.DataFrame() + jets_df['true_pt'] = true_pt + jets_df['pred_pt'] = pred_pt + jets_df['matched'] = 1 + jets_df = correct_jet_pt(jets_df, [0,5,10,20,30,50,100, 1000]) + relrespt = np.array((jets_df['true_pt'] - jets_df['pred_pt'])/ jets_df['true_pt']) + + #Start plotting + nParticle_bins = np.linspace(0, 30, 31) + fig_jet_nparticle = plt.figure() + plt.hist(nParicles_truth, bins=nParticle_bins, label='Truth', color='orange', alpha=0.5)#color='#fee7ae') + plt.grid(alpha=0.4) + plt.hist(nParicles_pred, bins=nParticle_bins, label='Prediction', histtype='step')#color='#67c4ce', linestyle='--') + plt.text(0.05, 0.95, datasetname, horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + plt.ylabel('Number of events', fontsize=20) + plt.xlabel('Jet nConstituents', fontsize=20) + plt.legend(fontsize=10) + plt.xticks(np.linspace(0, 30, 7)) + plt.close() + + relres_bins = np.linspace(-1, 1, 51) + fig_jet_relresE = plt.figure() + plt.hist(np.clip(relresE, -1, 1), bins=relres_bins, histtype='step')#color='#67c4ce', linestyle='--') + plt.grid(alpha=0.4) + write_mean_std_displayed(relresE, relres_bins, datasetname) + plt.ylabel('Number of events', fontsize=20) + plt.xlabel('rel. res. Jet Energy', fontsize=20) + plt.xticks(np.linspace(-1, 1, 5)) + plt.close() + + res_bins = np.linspace(-0.2, 0.2, 51) + fig_jet_reseta = plt.figure() + plt.hist(np.clip(reseta,-0.2,0.2), bins=res_bins, histtype='step')# color='#67c4ce', linestyle='--') + plt.grid(alpha=0.4) + write_mean_std_displayed(reseta, res_bins, datasetname) + plt.ylabel('Number of events', fontsize=20) + plt.xlabel('Jet $\\Delta \\eta$', fontsize=20) + plt.xticks(np.linspace(-0.2, 0.2, 5)) + plt.close() + + fig_jet_resphi = plt.figure() + plt.hist(np.clip(resphi, -0.2, 0.2), bins=res_bins, histtype='step')#color='#67c4ce', linestyle='--') + plt.grid(alpha=0.4) + write_mean_std_displayed(resphi, res_bins, datasetname) + plt.ylabel('Number of events', fontsize=20) + plt.xlabel('Jet $\\Delta \\phi$', fontsize=20) + plt.xticks(np.linspace(-0.2, 0.2, 5)) + plt.close() + + fig_jet_pt = plt.figure() + plt.hist(np.clip(relrespt, -1, 1), bins=relres_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(relrespt, relres_bins, datasetname) + plt.ylabel('Number of events', fontsize=20) + plt.xlabel('rel. res. Jet $p_T$', fontsize=20) + plt.xticks(np.linspace(-1, 1, 5)) + plt.close() + + return fig_jet_nparticle, fig_jet_relresE, fig_jet_reseta, fig_jet_resphi, fig_jet_pt + +def correct_jet_pt(jets_df, bins, plotEachBin=False, outputDir='/work/friemer/hgcalml/testplots/'): + """ + Function to calculate a correction factor for the jet pt in bins of pt, fit a correction function and apply it to the jets_df + """ + + corr_factors = [] + for i in range(len(bins) - 1): + mask_binpred = np.logical_and( + jets_df['pred_pt'] >= bins[i], + jets_df['pred_pt'] < bins[i + 1]) + mask = np.logical_and( + jets_df['matched'], + mask_binpred) + jet_pt_truth = jets_df[mask]['true_pt'] + jet_pt_pred = jets_df[mask]['pred_pt'] + respt = np.array((jet_pt_truth-jet_pt_pred)) + relrespt = np.array((jet_pt_truth-jet_pt_pred)/jet_pt_truth) + + if plotEachBin: + relres_bins = np.linspace(-1, 1, 51) + fig_jet_pt = plt.figure() + plt.hist(np.clip(relrespt, -1,1), bins=relres_bins, histtype='step') + plt.grid(alpha=0.4) + plt.text(0.05, 0.95, 'Quark Jet', horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + write_mean_std_displayed(relrespt, relres_bins) + plt.ylabel('Number of jets', fontsize=20) + plt.xlabel('rel. res. Cal. Jet $p_T$', fontsize=20) + plt.xticks(np.round(np.linspace(-1, 1, 9), 2), fontsize=12) + plt.savefig(os.path.join(outputDir, f'/jet_pt_{bins[i]}.png')) + plt.close() + + corr_factor = np.mean(relrespt) + corr_factors.append(corr_factor) + + #fit function to correction factor over bin middle + corr_factors = np.array(corr_factors) + + bins = np.array(bins) + bin_middle = bins[:-1] + (bins[1:] - bins[:-1])/2 + + mask = np.logical_and(~np.isnan(corr_factors), ~np.isinf(corr_factors)) #Should not be necessary, unless evaluation is performed on very small dataset + + def linearfunc(x, a, b): + return a*x+b + + popt, pcov = curve_fit(linearfunc, bin_middle[mask], corr_factors[mask]) + + corrected_jets_df = jets_df + corrected_jets_df['pred_pt'] = corrected_jets_df['pred_pt']*(1 + linearfunc(corrected_jets_df['pred_pt'], *popt)) + return corrected_jets_df + + +def plot_jet_metrics_from_jets(jets_df, datasetname='Quark Jet'): + """ + plots the jet metrics from the results of jet clustering algorithm + """ + reseta = [] + resphi = [] + relrespt = [] + nParicles_pred = [] + nParicles_truth = [] + + jets_df = correct_jet_pt(jets_df, [0,5,10,20,30,50,100, 1000]) + + for i in range(len(jets_df['event_id'].unique())): + matched = jets_df['matched'] + event_id_mask = jets_df['event_id'] == i + mask = np.logical_and(event_id_mask, matched) + + nParicles_truth.append(jets_df[mask]['true_n_constituents']) + nParicles_pred.append(jets_df[mask]['pred_n_constituents']) + + t_eta = jets_df[mask]['true_eta'] + pred_eta = jets_df[mask]['pred_eta'] + reseta.append(t_eta-pred_eta) + + t_phi = jets_df[mask]['true_phi'] + pred_phi = jets_df[mask]['pred_phi'] + delta_phi = calc_deltaphi(t_phi, pred_phi) + resphi.append(delta_phi) + + jet_pt_truth = jets_df[mask]['true_pt'] + jet_pt_pred = jets_df[mask]['pred_pt'] + relrespt.append((jet_pt_truth-jet_pt_pred)/jet_pt_truth) + + print('Total unmatched jets:', np.sum([jets_df['matched'] == False])) + + nParicles_truth = np.concatenate(nParicles_truth) + nParicles_pred = np.concatenate(nParicles_pred) + reseta = np.concatenate(reseta) + resphi = np.concatenate(resphi) + relrespt = np.concatenate(relrespt) + + nParticle_bins = np.linspace(0, 30, 31) + fig_jet_nparticle = plt.figure() + plt.hist(nParicles_truth, bins=nParticle_bins, label='Truth', color='orange', alpha=0.5) + plt.grid(alpha=0.4) + plt.hist(nParicles_pred, bins=nParticle_bins, label='Prediction', histtype='step') + plt.text(0.05, 0.95, datasetname, horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + plt.ylabel('Number of jets', fontsize=20) + plt.xlabel('Jet nConstituents', fontsize=20) + plt.legend(fontsize=10) + plt.xticks(np.linspace(0, 30, 7)) + plt.close() + + res_bins = np.linspace(-0.2, 0.2, 51) + fig_jet_reseta = plt.figure() + plt.hist(np.clip(reseta,-0.2,0.2), bins=res_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(reseta, res_bins, datasetname) + plt.ylabel('Number of jets', fontsize=20) + plt.xlabel('Jet $\\Delta \\eta$', fontsize=20) + plt.xticks(np.linspace(-0.2, 0.2, 5)) + plt.close() + + fig_jet_resphi = plt.figure() + plt.hist(np.clip(resphi, -0.2, 0.2), bins=res_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(resphi, res_bins, datasetname) + plt.ylabel('Number of jets', fontsize=20) + plt.xlabel('Jet $\\Delta \\phi$', fontsize=20) + plt.xticks(np.linspace(-0.2, 0.2, 5)) + plt.close() + + relres_bins = np.linspace(-1, 1, 51) + fig_jet_pt = plt.figure() + plt.hist(np.clip(relrespt, -1,1), bins=relres_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(relrespt, relres_bins, datasetname) + plt.ylabel('Number of jets', fontsize=20) + plt.xlabel('rel. res. Cal. Jet $p_T$', fontsize=20) + plt.xticks(np.round(np.linspace(-1, 1, 9), 2), fontsize=12) + plt.close() + + return fig_jet_nparticle, fig_jet_reseta, fig_jet_resphi, fig_jet_pt + + +def write_mean_std_displayed(data, bins, datasetname='Quark Jet'): + """ + Helper function to write the mean, std and displayed fraction of the data into the plot + """ + #remove inf and Nan + data = np.array(data, dtype=np.float64) + data = data[np.isfinite(data)] + data = data[np.isnan(data) == False] + + displayed = np.sum(np.logical_and(data >= bins[0], data < bins[-1]))/len(data) + mean= np.mean(data) + std = np.std(data) + + #Fit a normal distribution to the data + data = data[np.logical_and(data >= bins[0], data < bins[-1])] + mu, sigma = norm.fit(data) + + textstr = '\n'.join(( + # r'$\mathrm{Mean}=%.2f$' % (mean, ), + # r'$\mathrm{Std\ }=%.2f$' % (std, ), + r'$\mu=%.2f$' % (mu, ), + r'$\sigma=%.2f$' % (sigma, ), + r'$\mathrm{Displayed}=%.2f$' % (displayed, ))) + + # Add the text box + plt.text(0.95, 0.95, textstr, transform=plt.gca().transAxes, + fontsize=12, verticalalignment='top', horizontalalignment='right', + bbox=dict(boxstyle='round,pad=0.5', facecolor='white', edgecolor='black')) + plt.text(0.05, 0.95, datasetname, horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + + +def calc_matched_mask(df): + has_truth = np.isnan(df['truthHitAssignedEnergies']) == False + has_pred = np.isnan(df['pred_energy']) == False + matched = np.logical_and(has_truth, has_pred) + return matched + +def calc_deltaphi(t_phi, pred_phi): + delta_phi = t_phi-pred_phi + delta_phi = np.where(delta_phi > np.pi, delta_phi-2*np.pi, delta_phi) + delta_phi = np.where(delta_phi < -np.pi, delta_phi+2*np.pi, delta_phi) + return delta_phi + +def plot_particle_metrics(df, datasetname='Quark Jet'): + + matched = calc_matched_mask(df) + + reseta = [] + resphi = [] + relresE = [] + relrespt = [] + + for i in range(len(df['event_id'].unique())): + mask = np.logical_and(df['event_id'] == i, matched) + + truth_e = df[mask]['truthHitAssignedEnergies'] + pred_e = df[mask]['pred_energy'] + newrelresE = ((truth_e-pred_e)/truth_e).to_numpy() + relresE = np.concatenate((relresE, newrelresE)) + + t_eta = df[mask]['truthHitAssignedZ'] + pred_eta = df[mask]['pred_pos'].apply(lambda x: x[2]) + delta_eta = (t_eta-pred_eta).to_numpy() + reseta = np.concatenate((reseta, delta_eta)) + + t_phi = np.arctan2(df[mask]['truthHitAssignedY'], df[mask]['truthHitAssignedX']) + pred_phi = df[mask]['pred_pos'].apply(lambda x: np.arctan2(x[1], x[0])) + delta_phi = calc_deltaphi(t_phi, pred_phi) + resphi = np.concatenate((resphi, delta_phi)) + + truth_pt = truth_e/np.cosh(t_eta) + pred_pt = pred_e/np.cosh(pred_eta) + newrelrespt = ((truth_pt-pred_pt)/truth_pt).to_numpy() + relrespt = np.concatenate((relrespt, newrelrespt)) + + relres_bins = np.linspace(-1, 1, 51) + fig_relresE = plt.figure() + plt.hist(np.clip(relresE, -1,1), bins=relres_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(relresE, relres_bins, datasetname) + plt.xlabel('rel. res. Neutral Particle Energy', fontsize=20) + plt.ylabel('Number of neutral particles', fontsize=20) + plt.xticks(np.linspace(-1, 1, 5)) + plt.close() + + res_bins = np.linspace(-0.4, 0.4, 51) + fig_reseta = plt.figure() + plt.hist(np.clip(reseta, -0.4, 0.4), bins=res_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(reseta, res_bins, datasetname) + plt.xlabel('Neutral Particle $\\Delta \\eta$', fontsize=20) + plt.ylabel('Number of neutral particles', fontsize=20) + plt.xticks(np.linspace(-0.4, 0.4, 5)) + plt.close() + + fig_resphi = plt.figure() + plt.hist(np.clip(resphi, -0.4, 0.4), bins=res_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(resphi, res_bins, datasetname) + plt.xlabel('Neutral Particle $\\Delta \\phi$', fontsize=20) + plt.ylabel('Number of neutral particles', fontsize=20) + plt.xticks(np.linspace(-0.4, 0.4, 5)) + plt.close() + + fig_pt = plt.figure() + plt.hist(np.clip(relrespt, -1, 1), bins=relres_bins, histtype='step') + plt.grid(alpha=0.4) + write_mean_std_displayed(relrespt, relres_bins, datasetname) + plt.xlabel('rel. res. Neutral Particle $p_T$', fontsize=20) + plt.ylabel('Number of neutral particles', fontsize=20) + plt.xticks(np.linspace(-1, 1, 5)) + plt.close() + + return fig_relresE, fig_reseta, fig_resphi, fig_pt + +def calc_energy_resolution(df, bins=None): + matched = calc_matched_mask(df) + + ratios = [] #Relative deviation of predicted energy from truth energy + ratios_err= [] + diffs = [] #Absolute deviation of predicted energy from truth energy + diffs_err = [] + + for i in range(len(bins) - 1): + mask_truth_bin = np.logical_and( + df['truthHitAssignedEnergies'] >= bins[i], + df['truthHitAssignedEnergies'] < bins[i + 1]) + mask_bin = np.logical_and(mask_truth_bin, matched) + diff = np.abs(df['pred_energy'][mask_bin] - df['truthHitAssignedEnergies'][mask_bin]) + diffs.append(np.mean(diff)) + diffs_err.append(np.std(diff) / np.sqrt(len(diff))) + + ratio = df['pred_energy'][mask_bin] / df['truthHitAssignedEnergies'][mask_bin] + ratios.append(np.mean(ratio)) + ratios_err.append(np.std(ratio) / np.sqrt(len(ratio))) + + diffs, diffs_err, ratios, ratios_err = np.array(diffs), np.array(diffs_err), np.array(ratios), np.array(ratios_err) + return diffs, diffs_err, ratios, ratios_err +def plot_energy_resolution(df, bins=None, datasetname='Quark Jet'): + diffs, diffs_err, ratios, ratios_err = calc_energy_resolution(df, bins) + + neutral_mask = df['truthHitAssignedPIDs'].isin([22,130,310,311,2112,-2112,3122,-3122,3322,-3322]) + diffs_neutral, diffs_err_neutral, ratios_neutral, ratios_err_neutral = calc_energy_resolution(df[neutral_mask], bins) + + charged_mask = df['truthHitAssignedPIDs'].isin([11,-11,13,-13,211,-211,321,-321,2212,-2212,3112,-3112,3222,-3222,3312,-3312]) + diffs_charged, diffs_err_charged, ratios_charged, ratios_err_charged = calc_energy_resolution(df[charged_mask], bins) + + # Calculate the bin positions and widths + binwidth = bins[1:] - bins[:-1] + x_pos = bins[:-1] + binwidth / 2 + x_err = binwidth / 2 + # Create the stacked plots of diff and ratio with one errorbar for total, neutral and charged each + fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(12, 8), sharex=True, gridspec_kw={'height_ratios': [3, 1]}) + + ax1.errorbar(x_pos, diffs, xerr=x_err, yerr=diffs_err, fmt='o', label='Total') + ax1.errorbar(x_pos, diffs_neutral, xerr=x_err, yerr=diffs_err_neutral, fmt='o', label='Neutral') + ax1.errorbar(x_pos, diffs_charged, xerr=x_err, yerr=diffs_err_charged, fmt='o', label='Charged') + ax1.set_ylabel('$\Delta p \ [GeV]$') + ax1.grid(alpha=0.4) + #write datasetname in top right corner + ax1.text(0.05, 0.65, datasetname, horizontalalignment='left', verticalalignment='top', transform=ax1.transAxes, fontsize=10) + ax1.legend() + + ax2.errorbar(x_pos, ratios, xerr=x_err, yerr=ratios_err, fmt='o', label='Total') + ax2.errorbar(x_pos, ratios_neutral, xerr=x_err, yerr=ratios_err_neutral, fmt='o', label='Neutral') + ax2.errorbar(x_pos, ratios_charged, xerr=x_err, yerr=ratios_err_charged, fmt='o', label='Charged') + ax2.set_ylabel('$p_{pred} / p_{truth}$') + ax2.grid(alpha=0.4) + + ydelta = np.round(max(1 - ax2.get_ylim()[0], ax2.get_ylim()[1] - 1),1) + ax2.set_ylim(1 - ydelta, 1 + ydelta) + + ax2.set_xlim(bins[0], bins[-1]) + ax2.set_xscale('log') + ax2.set_xlabel('Truth Energy [GeV]') + ax2.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter()) + xticks = np.round(bins, 0).astype(int) + ax2.set_xticks(xticks) + ax2.set_xticklabels(xticks) + + fig.subplots_adjust(hspace=0.1) + + return fig + +def plot_condensation(pred_dict, t_dict, feature_dict, coordspace='condensation'): + """ + Plots the condensation space of a event (coordspace='condensation") or the input space (coordspace='input') + """ + + if(coordspace == 'condensation'): + coords = pred_dict['pred_ccoords'] + + if coords.shape[1] < 3: #add a zero + coords = tf.concat([coords, tf.zeros_like(coords)[:,0:1]], axis=-1) + X = coords[:,0] + Y = coords[:,1] + Z = coords[:,2] + sizefield = 'beta' + elif(coordspace == 'input'): + X = feature_dict['recHitX'][:,0] + Y = feature_dict['recHitY'][:,0] + Z = feature_dict['recHitZ'][:,0] + sizefield = 'rechit_energy_scaled' + else: + raise ValueError('coordspace must be "condensation" or "input"') + + data={ + 'X': X, + 'Y': Y, + 'Z': Z, + 't_idx': t_dict['truthHitAssignementIdx'][:,0], + 'beta': pred_dict['pred_beta'][:,0], + 'pdgid': t_dict['truthHitAssignedPIDs'][:,0], + 'recHitID': feature_dict['recHitID'][:,0], + 'rechit_energy': feature_dict['recHitEnergy'][:,0], + 'rechit_energy_scaled': np.log(feature_dict['recHitEnergy'][:,0]+1), + } + eventdf = pd.DataFrame(data) + + hover_data = {'X': True, 'Y': True, 'Z': True, 't_idx': True,'beta': True, 'pdgid': True,'recHitID': True, 'rechit_energy': True} + + fig = px.scatter_3d(eventdf, x="X", y="Y", z="Z", + color="t_idx", + size=sizefield, + template='plotly_white', + hover_data=hover_data, + color_continuous_scale=px.colors.sequential.Rainbow) + fig.update_traces(marker=dict(line=dict(width=0))) + + return fig + +def plot_distribution(truth_list, feature_list): + """ + Makes boxplots of the number of particles, tracks, hits, etc. in the events + """ + n_tracks = [] + n_hits = [] + n_particles = [] + n_charged_had = [] + n_lepton = [] + n_photon = [] + n_neutral_had = [] + for i in range(len(feature_list)): + n_tracks.append(np.sum(np.abs(feature_list[i]['recHitID']))) + n_hits.append(len(feature_list[i]['recHitID'])-n_tracks[-1]) + truth_data ={ + 'truthHitAssignementIdx': truth_list[i]['truthHitAssignementIdx'][:,0], + 'truthHitAssignedPIDs': truth_list[i]['truthHitAssignedPIDs'][:,0], + } + truth_df = pd.DataFrame(truth_data) + n_particles.append(len(truth_df['truthHitAssignementIdx'].unique())) + + charged_hd_mask = truth_df['truthHitAssignedPIDs'].isin([211,-211,321,-321,2212,-2212,3112,-3112,3222,-3222,3312,-3312]) + lepton_mask = truth_df['truthHitAssignedPIDs'].isin([11,-11,13,-13]) + photon_mask = truth_df['truthHitAssignedPIDs'].isin([22]) + neutral_mask = truth_df['truthHitAssignedPIDs'].isin([130,310,311,2112,-2112,3122,-3122,3322,-3322]) + n_charged_had.append(len(truth_df[charged_hd_mask]['truthHitAssignementIdx'].unique())) + n_lepton.append(len(truth_df[lepton_mask]['truthHitAssignementIdx'].unique())) + n_photon.append(len(truth_df[photon_mask]['truthHitAssignementIdx'].unique())) + n_neutral_had.append(len(truth_df[neutral_mask]['truthHitAssignementIdx'].unique())) + #convert to numpy + n_tracks, n_hits, n_particles, n_charged_had, n_lepton, n_photon, n_neutral_had = \ + np.array(n_tracks), np.array(n_hits), np.array(n_particles), np.array(n_charged_had), np.array(n_lepton), np.array(n_photon), np.array(n_neutral_had) + + #create horizontal boxplot + fig = plt.figure() + values = [n_lepton, n_charged_had, n_neutral_had, n_photon, n_particles, n_tracks,n_hits/100] + labels = ['leptons', 'ch. hadrons','nu. hadrons','photons','total particles', 'tracks','cells [$10^2$]'] + + plt.boxplot(values, labels=labels, vert=False, patch_artist=True, + boxprops=dict(facecolor="limegreen"), medianprops=dict(color="black"), + whis = (0, 100)) + + plt.grid(alpha=0.4, axis='x') + # Add mean values as text + for pos, data in zip(np.arange(len(values)), values): + mean = np.mean(data) + plt.text(21.2, pos+1, f'({mean:.1f})', va='center', ha='left') + + plt.close() + return fig +def plot_truth_jet_nconstrituents(jets_df, df, datasetname='Quark Jet'): + nParicles_truth_jet = [] + + jets_df = correct_jet_pt(jets_df, [0,5,10,20,30,50,100, 1000]) + matched = jets_df['matched'] + for i in range(len(jets_df['event_id'].unique())): + event_id_mask = jets_df['event_id'] == i + mask = np.logical_and(event_id_mask, matched) + + nParicles_truth_jet.append(jets_df[mask]['true_n_constituents']) + nParicles_truth_jet = np.concatenate(nParicles_truth_jet) + + nParicles_truth_particle = [] + has_truth = np.isnan(df['truthHitAssignedEnergies']) == False + + for i in range(len(df['event_id'].unique())): + event_mask = df['event_id'] == i + truth_mask = np.logical_and(event_mask, has_truth) + nParicles_truth_particle.append(np.sum(truth_mask)) + nParicles_truth_particle = np.array(nParicles_truth_particle) + + nParticle_bins = np.linspace(0, 30, 31) + fig_nparticle_truth_compare = plt.figure() + plt.hist(nParicles_truth_jet, bins=nParticle_bins, label='Anti-$k_T$ based', color='blue', alpha=0.5) + plt.hist(nParicles_truth_particle, bins=nParticle_bins, label='Event based', color='orange', alpha=0.5) + plt.grid(alpha=0.4) + plt.text(0.05, 0.95, datasetname, horizontalalignment='left', verticalalignment='top', transform=plt.gca().transAxes, fontsize=10) + plt.ylabel('Number of jets', fontsize=20) + plt.xlabel('Jet nConstituents', fontsize=20) + plt.legend(fontsize=12) + plt.xticks(np.linspace(0, 30, 7)) + plt.close() + return fig_nparticle_truth_compare + +def plot_energydistribution(df): + fig, ax = plt.subplots(figsize=(12, 6)) + ax.hist(df['truthHitAssignedEnergies'], histtype='step', color='blue', label='Truth') + ax.hist(df['pred_energy'], histtype='step', label='Prediction') + ax.set_yscale('log') + ax.set_xscale('log') + ax.set_xlabel('Energy [GeV]') + ax.set_ylabel('Number of showers') + ax.legend() + ax.grid(alpha=0.4) + return fig + + +def plot_everything(df, pred_list, t_list, feature_list, jets_df, outputpath='/work/friemer/hgcalml/testplots/', datasetname='Quark Jet'): + + #Create Output directory + if not os.path.exists(outputpath): + print("Creating output directory", outputpath) + os.mkdir(outputpath) + else: + var = input(\ + 'Working directory exists, are you sure you want to continue, please type "yes/y"\n') + var = var.lower() + if not var in ('yes', 'y'): + sys.exit() + os.makedirs(os.path.join(outputpath, 'condensation'), exist_ok=True) + + matplotlib.rcParams.update({'font.size': 20}) + + print('Plotting efficiencies') + energy_bins_neutral = np.array([1,2,3,4,5,10,20,30,50]) + fig_eff_overview=plot_efficencies(df, bins=energy_bins_neutral) + fig_eff_overview.savefig(os.path.join(outputpath,'efficienciesOverview_{}.png'.format(datasetname)), bbox_inches='tight') + + fig_eff, fig_fake = plot_efficency_and_fakerate(df, bins=energy_bins_neutral, datasetname=datasetname) + fig_eff.savefig(os.path.join(outputpath,'efficiencies_{}.png'.format(datasetname)), bbox_inches='tight') + fig_fake.savefig(os.path.join(outputpath,'fake_rate_{}.png'.format(datasetname)), bbox_inches='tight') + + print('Plotting jet metrics') + fig_jet_nparticle_all, fig_jet_relresE_all, fig_jet_reseta_all, fig_jet_resphi_all, fig_jet_pt_all = plot_jet_metrics_based_on_particles(df, datasetname=datasetname) + fig_jet_nparticle_all.savefig(os.path.join(outputpath,'jet_nparticle_all_{}.png'.format(datasetname)), bbox_inches='tight') + fig_jet_relresE_all.savefig(os.path.join(outputpath,'jet_relresE_all_{}.png'.format(datasetname)), bbox_inches='tight') + fig_jet_reseta_all.savefig(os.path.join(outputpath,'jet_reseta_all_{}.png'.format(datasetname)), bbox_inches='tight') + fig_jet_resphi_all.savefig(os.path.join(outputpath,'jet_resphi_all_{}.png'.format(datasetname)), bbox_inches='tight') + fig_jet_pt_all.savefig(os.path.join(outputpath,'jet_pt_all_{}.png'.format(datasetname)), bbox_inches='tight') + + fig_jet_nparticle_clustered, fig_jet_reseta_clustered, fig_jet_resphi_clustered, fig_jet_pt_clustered = plot_jet_metrics_from_jets(jets_df, datasetname=datasetname) + fig_jet_nparticle_clustered.savefig(os.path.join(outputpath,'jet_nparticle_clustered_{}.png'.format(datasetname)), bbox_inches='tight') + fig_jet_reseta_clustered.savefig(os.path.join(outputpath,'jet_reseta_clustered_{}.png'.format(datasetname)), bbox_inches='tight') + fig_jet_resphi_clustered.savefig(os.path.join(outputpath,'jet_resphi_clustered_{}.png'.format(datasetname)), bbox_inches='tight') + fig_jet_pt_clustered.savefig(os.path.join(outputpath,'jet_pt_clustered_{}.png'.format(datasetname)), bbox_inches='tight') + + print('Plotting particle metrics') + mask_neutral = df['truthHitAssignedPIDs'].isin([22,130, 310, 311, 2112, -2112, 3122, -3122, 3322, -3322]) + fig_neutral_relresE, fig_neutral_reseta, fig_neutral_resphi, fig_neutral_pt = plot_particle_metrics(df[mask_neutral], datasetname=datasetname) + fig_neutral_relresE.savefig(os.path.join(outputpath,'neutral_relresE_{}.png'.format(datasetname)), bbox_inches='tight') + fig_neutral_reseta.savefig(os.path.join(outputpath,'neutral_reseta_{}.png'.format(datasetname)), bbox_inches='tight') + fig_neutral_resphi.savefig(os.path.join(outputpath,'neutral_resphi_{}.png'.format(datasetname)), bbox_inches='tight') + fig_neutral_pt.savefig(os.path.join(outputpath,'neutral_pt_{}.png'.format(datasetname)), bbox_inches='tight') + + fig_class = plot_classification_p(df, energy_bins_neutral, datasetname=datasetname) + fig_class.savefig(os.path.join(outputpath,'classification_p_{}.png'.format(datasetname)), bbox_inches='tight') + + print('Plotting energy resolution') + energy_bins_combined = np.array([1,5,10,20,30,50,200]) + fig_new_res = plot_energy_resolution(df, bins=energy_bins_combined, datasetname=datasetname) + fig_new_res.savefig(os.path.join(outputpath,'energy_resolution_{}.png'.format(datasetname)), bbox_inches='tight') + + print('Plotting condensation and input') + for event_id in range(10): + fig = plot_condensation(pred_list[event_id], t_list[event_id], feature_list[event_id], 'condensation') + fig.write_html(os.path.join(outputpath, 'condensation' ,'condensation{}_{}.html'.format(event_id, datasetname))) + fig = plot_condensation(pred_list[event_id], t_list[event_id], feature_list[event_id], 'input') + fig.write_html(os.path.join(outputpath, 'condensation' ,'input{}_{}.html'.format(event_id, datasetname))) + + # print('Plotting truth compare') + # fig_nparticle_truth_compare = plot_truth_jet_nconstrituents(jets_df, df, datasetname=datasetname) + # fig_nparticle_truth_compare.savefig(os.path.join(outputpath,'nparticle_truth_compare_{}.png'.format(datasetname)), bbox_inches='tight') + + # print('Plotting distribution') + # fig_dist = plot_distribution(t_list, feature_list) + # fig_dist.savefig(os.path.join(outputpath, 'distribution_{}.png'.format(datasetname)), bbox_inches='tight') + +def plot_everything_from_file(analysisfilepath, outputpath='/work/friemer/hgcalml/testplots/', datasetname='Quark Jet'): + #Load analysis file created by analyse_cocoa_predictions.py + with gzip.open(analysisfilepath, 'rb') as input_file: + analysis_data = pickle.load(input_file) + + df = analysis_data['showers_dataframe'] + pred_list = analysis_data['prediction'] + t_list = analysis_data['truth'] + feature_list = analysis_data['features'] + jets_df = analysis_data['jets_dataframe'] + + plot_everything(df, pred_list, t_list, feature_list, jets_df, outputpath, datasetname=datasetname) + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + 'Create plots for the COCOA analysis') + parser.add_argument('analysisfile', + help='Filepath to analysis file created by analyse_cocoa_predictions.py containing shower dataframe') + parser.add_argument('outputlocation', + help="Output directory for the plots", + default='') + parser.add_argument('--gluon', action='store_true', help='Write Gluon dataset instead of Quark dataset in top left corner') + args = parser.parse_args() + + if args.gluon: + datasetname='Gluon Jet' + else: + datasetname='Quark Jet' + plot_everything_from_file(args.analysisfile, args.outputlocation, datasetname=datasetname) \ No newline at end of file diff --git a/scripts/submitViaHTCondor.py b/scripts/submitViaHTCondor.py new file mode 100644 index 00000000..9ef59ef5 --- /dev/null +++ b/scripts/submitViaHTCondor.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python3 +""" +Script to submit training scripts to Topas via HTCondor +""" +# Usage: +# python3 submitViaHTCondor.py ---n JobName +# +# Example usage for a python script with 3 arguments +# python3 submitViaHTCondor python3 mypythonscript.py --arg1 --arg2 --arg3 ---n JobName + + +#Imports +import sys +import os +import uuid +import subprocess + +#Usefull class to handle meta options +class meta_option(object): + def __init__(self, id, default_val = None) -> None: + self.id = id + self.triggered = False + self.value = default_val + + def check(self, clo): + if self.triggered: + self.value = clo + self.triggered = False + return True + if clo == '---'+self.id: + self.triggered = True + return True + return False + + def valid(self): + return self.value is not None + + def __str__(self) -> str: + return self.id + ': '+self.value + +UEXT = str(uuid.uuid4()) +#Define possible meta parameters +opts = { + 'n' : meta_option('n','TrainJob_'+UEXT), + 'f' : meta_option('f', '/work/friemer/hgcalml/HGCalML'), + 'cpu': meta_option('cpu', '1'), + 'memory': meta_option('memory', '20 GB'), + 'disk': meta_option('disk', '10 GB'), + 'gpu': meta_option('gpu', '1'), + 'wandb': meta_option('wandb', 'True'), + 'submit': meta_option('submit', 'True') +} + +#Filter out meta parameters +filtered_clo=[] + +for clo in sys.argv: + next = False + for _,o in opts.items(): + if o.check(clo): + next = True + break + if next: + continue + filtered_clo.append(clo) + +all_valid = True +for _,o in opts.items(): + all_valid = all_valid and o.valid() + +filtered_clo = filtered_clo[1:] +COMMANDS = " " +for clos in filtered_clo: + COMMANDS += clos + " " + +#convert strings to bool +opts['wandb'].value = opts['wandb'].value.lower() in ('true', 't', 'yes', 'y','1') +opts['submit'].value = opts['submit'].value.lower() in ('true', 't', 'yes', 'y', '1') + +#Help messages +if '-h' in sys.argv or '--help' in sys.argv or (not all_valid): + print('script to create a submission and script file and submit them via HTCondor to topas\n') + print('all commands are fully forwarded with one exception:') + print('\n ---n (opt) specifies a name for the scripts\n') + print('\n ---f location of folder with other necessary files\n') + print('\n ---cpu (opt) number of cpus to request default: 1\n') + print('\n ---memory (opt) size of memory to request default: 15 GB\n') + print('\n ---disk (opt) size of memory to request default 8 GB\n') + print('\n ---gpu (opt) number of gpus to request default: 1\n') + print('\n ---wandb (opt) search for a Wandb API Key default: True\n') + print('\n ---submit (opt) submit the files in the end default: True\n') + sys.exit() + +#Create dictionary for files +if os.path.isdir(opts['n'].value): + var = input(\ + 'Working directory exists, are you sure you want to continue, please type "yes/y"\n') + var = var.lower() + if not var in ('yes', 'y'): + sys.exit() +else: + os.system('mkdir -p '+opts['n'].value) + + +#Get current working directory +CWD = os.getcwd() + + +#################################################################################################### +### Create Sub File ################################################################################ +#################################################################################################### + +#Get absolute filepath and transfer entire folder if there is a .djcdc file also replace filepaths with filenames in command +inputfileslocations ='' +NEWCOMMANDS = '' +for word in COMMANDS.split(): + if os.path.exists(word): + if '.djcdc' in word: + inputfileslocations += os.path.join(CWD, os.path.dirname(word))+ '/, ' + else: + inputfileslocations += os.path.join(CWD, word) + ', ' + NEWCOMMANDS+=os.path.basename(word) + ' ' + else: + NEWCOMMANDS+=word + ' ' + +#Create a tarball of the folder +hasfolder = True +if os.path.isdir(opts['f'].value): + print('Folder found, creating a tarball of the folder...') + os.system(f''' + cd {opts['f'].value} + cd ../ + tar -czf {opts['n'].value}/ZipFolder.tar.gz {os.path.basename(opts['f'].value)} + cd {CWD}''') + inputfileslocations += CWD+ '/' +opts['n'].value + '/ZipFolder.tar.gz, ' +elif os.path.isfile(opts['f'].value) and opts['f'].value.endswith('.tar.gz'): + print('File found, copying the file to the submission folder...') + os.system(f''' + cp {opts['f'].value} {opts['n'].value}/ZipFolder.tar.gz''') + inputfileslocations += CWD+ '/' +opts['n'].value + '/ZipFolder.tar.gz, ' +else: + print('Folder not found, no additional files will be transferred.') + hasfolder = False + + +#Setup the sub file +print('Creating the submission file...') +sub_temp=f'''#!/bin/bash +Universe = docker +docker_image = cernml4reco/deepjetcore4:abc9aee +accounting_group = cms.jet + +requirements =(TARGET.CloudSite=="topas") ++RemoteJob = True ++RequestWalltime = 24 * 60 * 60 + +executable = {opts['n'].value+'_run.sh'} + +should_transfer_files = YES +when_to_transfer_output = ON_EXIT + +transfer_input_files = {inputfileslocations} {CWD+ '/' +opts['n'].value+ '/' + opts['n'].value+'_run.sh'} +transfer_output_files = . + +output = {opts['n'].value}.out +error = {opts['n'].value}.err +log = {opts['n'].value}.log + +request_cpus = {opts['cpu'].value} +request_memory = {opts['memory'].value} +request_disk = {opts['disk'].value} +request_GPUs = {opts['gpu'].value} +''' + +#Get Wandb API Key +if opts['wandb'].value: + # Command to get WANDB_API_KEY + script_path = os.path.expanduser('~/private/wandb_api.sh') + command = f''' + echo "Checking for wandb API Key" + if [ -f {script_path} ]; then + source {script_path} + echo $WANDB_API_KEY + fi''' + result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, executable='/bin/bash') + + # Extract the WANDB_API_KEY from the command output + API_Key = result.stdout.strip().split('\n')[-1] + if API_Key is not None: + sub_temp += f'''environment = "WANDB_API_KEY={API_Key}" + ''' + else: + print('No WANDB_API_KEY found, please set it in your environment or in ~/private/wandb_api.sh or use the --no_wandb flag.') + + +#End of the sub file +sub_temp += f''' +queue +''' + +with open(opts['n'].value+'/'+opts['n'].value+'.sub','w', encoding='utf-8') as f: + f.write(sub_temp) + + +#################################################################################################### +### Create Run File ################################################################################ +#################################################################################################### + +print('Creating the run file...') +#Unpack the ZipFolder.tar.gz +runscript_temp=f''' +#!/bin/bash +''' + +if hasfolder: + runscript_temp += f''' +#Unpack the ZipFolder.tar.gz and set the necessary environment variables +tar -xzf ZipFolder.tar.gz + +#Display the content of the directory for debugging +ls -l + +''' + +#if the folder is my HGCalML, set the necessary environment variables +if os.path.basename(opts['f'].value)=='HGCalML': + runscript_temp += f''' +#Set some environment variables +export HGCALML=$(readlink -f HGCalML) +export DEEPJETCORE_SUBPACKAGE=$HGCALML +export PATH=$HGCALML/scripts:$PATH +export PYTHONPATH=$HGCALML/modules:$PYTHONPATH +export PYTHONPATH=$HGCALML/scripts:$PYTHONPATH +export LD_LIBRARY_PATH=$HGCALML/modules:$LD_LIBRARY_PATH +export LC_ALL=C.UTF-8 # necessary for wandb +export LANG=C.UTF-8 # necessary for wandb +''' + +#Run the commands +runscript_temp += f''' +#Run the command +{NEWCOMMANDS} +''' + +#if the folder is my HGCalML, clean up unnecessary files +if os.path.basename(opts['f'].value)=='HGCalML': + runscript_temp += f''' +#Clean up +rm ZipFolder.tar.gz +rm -r HGCalML +rm *.djctd +rm *.djcdc +rm -r tmp +rm -r var +rm -r wandb + +#Display the content of the directory for debugging +ls -l +''' +with open(opts['n'].value+'/'+opts['n'].value+'_run.sh','w', encoding='utf-8') as f: + f.write(runscript_temp) + + +#################################################################################################### +### Submit the Job ################################################################################# +#################################################################################################### +if opts['submit'].value: + SUBMITCOMMAND = ( + 'cd ' + opts['n'].value + '; pwd; condor_submit ' + opts['n'].value + '.sub' + ) + os.system(SUBMITCOMMAND) \ No newline at end of file diff --git a/tests/checkDJC.py b/tests/checkDJC.py new file mode 100644 index 00000000..e7f4422f --- /dev/null +++ b/tests/checkDJC.py @@ -0,0 +1,55 @@ +from datastructures import TrainData_Cocoa +from djcdata.dataPipeline import TrainDataGenerator +import numpy as np +import matplotlib.pyplot as plt + +td = TrainData_Cocoa() + +td.readFromFile('/work/friemer/hgcalml/trainingdata/singleQuarkJet_train.djctd') +gen = TrainDataGenerator() +gen.setBatchSize(1) +gen.setBuffer(td) +gen.setSkipTooLargeBatches(False) + +data = next(gen.feedTrainData())#this is a dict, row splits can be ignored, this is per event +df= TrainData_Cocoa.createPandasDataFrame(data, 0) +print(df.keys()) +# Index(['recHitEnergy', 'recHitEta', 'recHitHitR', 'recHitID', +# 'recHitLogEnergy', 'recHitR', 'recHitTheta', 'recHitTime', 'recHitX', +# 'recHitY', 'recHitZ', 't_rec_energy', 'truthHitAssignedEnergies', +# 'truthHitAssignedEta', 'truthHitAssignedPIDs', 'truthHitAssignedPhi', +# 'truthHitAssignedT', 'truthHitAssignedX', 'truthHitAssignedY', +# 'truthHitAssignedZ', 'truthHitAssignementIdx', +# 'truthHitFullyContainedFlag', 'truthHitSpectatorFlag'], +# dtype='object') + + +phi_truth = np.arctan2(df['truthHitAssignedY'], df['truthHitAssignedX']) +eta_truth = df['truthHitAssignedZ'] + +# phi_truth = df['truthHitAssignedPhi'] +# eta_truth = df['truthHitAssignedEta'] #Wrong definition of eta in the nanoML file vs cococa + +# print(max(abs(phi_pred - phi_truth))) +# print(max(abs(eta_pred - eta_truth))) + +phi_hit = np.arctan2(df['recHitY'], df['recHitX']) +#Calculate Eta from X,Y,Z +eta_hit = np.arcsinh(df['recHitZ']/np.sqrt(df['recHitX']**2 + df['recHitY']**2)) + +print(np.std(phi_hit - phi_truth)) +print(max(abs(phi_hit - phi_truth))) +plt.hist(phi_hit - phi_truth, bins=25) +plt.ylabel('number of hits') +plt.xlabel('phi_hit - phi_truth') +plt.savefig('phi_hit.png') + +print(np.std(eta_hit - eta_truth)) +print(max(abs(eta_hit - eta_truth))) +plt.hist(eta_hit - eta_truth, bins=25) +plt.ylabel('number of hits') +plt.xlabel('eta_hit - eta_truth') +plt.savefig('eta_hit.png') + + + diff --git a/tests/testOC.py b/tests/testOC.py new file mode 100644 index 00000000..68473cf1 --- /dev/null +++ b/tests/testOC.py @@ -0,0 +1,83 @@ +import tensorflow as tf +import numpy as np +from LossLayers import LLFullObjectCondensation +import random + +#Small test to check if the loss function is working as expected afer I changed the oc-Kernals + +def random_loss(seed): + tf.random.set_seed(seed) + random.seed(seed) + np.random.seed(seed) + + llFOBC = LLFullObjectCondensation(scale=1., + use_energy_weights=True, + record_metrics=True, + print_loss=True, + name="ExtendedOCLoss", + implementation = "hinge", + beta_loss_scale = 1.0, + too_much_beta_scale = 0.0, + energy_loss_weight = 0.4, + classification_loss_weight = -0.4, + position_loss_weight = 0.0, + timing_loss_weight = 0.0, + q_min = 1.0, + use_average_cc_pos = 0.9999) + + # Adjusting the inputs to have random values but similar to the provided examples + length = np.random.randint(1000, 1100) + + + pred_beta = tf.random.uniform((length, 1), minval=0.996, maxval=1.0, dtype=tf.float32) + pred_ccoords = tf.random.uniform((length, 3), minval=-76.6, maxval=292.2, dtype=tf.float32) + pred_distscale = tf.ones((length, 1), dtype=tf.float32) + pred_energy = tf.random.uniform((length, 1), minval=1.009, maxval=1.194, dtype=tf.float32) + pred_energy_low_quantile = tf.zeros((length, 1), dtype=tf.float32) + pred_energy_high_quantile = tf.random.uniform((length, 1), minval=0, maxval=165, dtype=tf.float32) + pred_pos = tf.random.uniform((length, 2), minval=-29.99, maxval=166.66, dtype=tf.float32) + pred_time = tf.random.uniform((length, 1), minval=7.211, maxval=10.004, dtype=tf.float32) + pred_time_unc = tf.random.uniform((length, 1), minval=0.974, maxval=3.322, dtype=tf.float32) + pred_id = tf.random.uniform((length, 6), minval=0, maxval=8, dtype=tf.float32) + rechit_energy = tf.random.uniform((length, 1), minval=13.72, maxval=4103.71, dtype=tf.float32) + t_idx = tf.random.uniform((length, 1), minval=1, maxval=8, dtype=tf.int32) + t_energy = tf.random.uniform((length, 1), minval=2302.9, maxval=6639.13, dtype=tf.float32) + t_pos = tf.random.uniform((length, 3), minval=-0.0027, maxval=0.0011, dtype=tf.float32) + t_time = tf.zeros((length, 1), dtype=tf.float32) + t_pid = tf.random.uniform((length, 1), minval=-211, maxval=211, dtype=tf.int32) + t_spectator_weights = tf.zeros((length, 1), dtype=tf.float32) + t_fully_contained = tf.ones((length, 1), dtype=tf.int32) + t_rec_energy = tf.random.uniform((length, 1), minval=812.48, maxval=2587.7, dtype=tf.float32) + t_is_unique = tf.random.uniform((length, 1), minval=0, maxval=1, dtype=tf.int32) + + #randomly generate rowsplits for testing + # Generate three random integers between 1 and length-1, ensuring they are sorted + num_splits = np.random.randint(1, 3) + random_splits = sorted(np.random.randint(8, length, size=num_splits)) + + # Append 0 at the beginning and 'length' at the end to form valid row splits + rowsplits = [0] + random_splits + [length] + + #rowsplits = [0] + [length] + + # Convert to a TensorFlow constant + rowsplits = tf.constant(rowsplits, dtype=tf.int32) + + + # Grouping inputs + inputs = [pred_beta, pred_ccoords, pred_distscale, + pred_energy, pred_energy_low_quantile, pred_energy_high_quantile, + pred_pos, pred_time, pred_time_unc, pred_id, + rechit_energy, + t_idx, t_energy, t_pos, t_time, t_pid, t_spectator_weights, t_fully_contained, t_rec_energy, + t_is_unique, + rowsplits] + + + loss = llFOBC.loss(inputs) # Adjust this call based on how the loss function is actually used + print(loss) + return loss + + +for i in [42,300,500,1,565656,400432,342342,3242342]: + random_loss(i) \ No newline at end of file diff --git a/tests/testloadmodel.py b/tests/testloadmodel.py new file mode 100644 index 00000000..372c185a --- /dev/null +++ b/tests/testloadmodel.py @@ -0,0 +1,57 @@ +from DeepJetCore.modeltools import load_model + +from DeepJetCore.DJCLosses import * +from DeepJetCore.DJCLayers import * +import imp +from Layers import * +from Losses import * +from Metrics import * +import h5py + + +def get_custom_objects(): + imp.find_module('Layers') + imp.find_module('Losses') + imp.find_module('Metrics') + custom_objs = {} + custom_objs.update(djc_global_loss_list) + custom_objs.update(djc_global_layers_list) + custom_objs.update(global_loss_list) + custom_objs.update(global_layers_list) + custom_objs.update(global_metrics_list) + return custom_objs + + +custom_objs = get_custom_objects() + +model_path = '/work/friemer/hgcalml/paper_BS100000__noDS_qm01_lr0005/model_results/KERAS_model.h5' + + +model_path = '/work/friemer/hgcalml/testCocoaLoad2/KERAS_model.h5' +# #Print shapes of all layers in model +# print('Model Path:', model_path) +# with h5py.File(model_path, 'r') as f: +# f.visititems(lambda name, obj: print(name, obj.shape if hasattr(obj, 'shape') else '')) + + + +# import h5py + +# with h5py.File(model_path, 'r') as f: +# f.visititems(lambda name, obj: print(name, obj.shape if hasattr(obj, 'shape') else '')) + + +model = load_model(model_path) + + + +# #print shapes of all layers in model and shapes of all weights +# for layer in model.layers: +# print(layer.name) +# print('input_shape:', layer.input_shape) +# print('output_shape',layer.output_shape) +# #print shape of weights +# print('shape of weights: ',[w.shape for w in layer.get_weights()]) +# print('') +model.save('testsavedMOdel.h5') +model2 = load_model('testsavedMOdel.h5') \ No newline at end of file