umami.train_tools package#

Submodules#

umami.train_tools.configuration module#

Configuration module for NN trainings.

class umami.train_tools.configuration.EvaluationSettingsConfig(evaluate_traind_model: bool = True, results_filename_extension: str = '', n_jets: int | None = None, tagger: list = <factory>, frac_values_comp: dict = <factory>, frac_values: dict = <factory>, working_point: float | None = None, eff_min: float = 0.49, eff_max: float = 1.0, frac_step: float = 0.01, frac_min: float = 0.01, frac_max: float = 1.0, add_eval_variables: list = <factory>, eval_batch_size: int | None = None, extra_classes_to_evaluate: list = <factory>, shapley: dict = <factory>, calculate_saliency: bool = False, saliency_ntrks: int = 8, saliency_effs: list | None = None, x_axis_granularity: int = 100, figsize: list = <factory>)#

Bases: object

Dataclass for the evaluation_settings config options.

add_eval_variables: list#

calculate_saliency: bool = False#

eff_max: float = 1.0#

eff_min: float = 0.49#

eval_batch_size: int = None#

evaluate_traind_model: bool = True#

extra_classes_to_evaluate: list#

figsize: list#

frac_max: float = 1.0#

frac_min: float = 0.01#

frac_step: float = 0.01#

frac_values: dict#

frac_values_comp: dict#

n_jets: int = None#

results_filename_extension: str = ''#

saliency_effs: list = None#

saliency_ntrks: int = 8#

shapley: dict#

tagger: list#

working_point: float = None#

x_axis_granularity: int = 100#

class umami.train_tools.configuration.NNStructureConfig(evaluate_trained_model: bool = True, class_labels: list | None = None, tagger: str | None = None, batch_size: int | None = None, learning_rate: float | None = None, epochs: int | None = None, main_class: str | None = None, n_jets_train: int | None = None, load_optimiser: bool = False, use_sample_weights: bool = False, nfiles_tfrecord: int = 5, batch_normalisation: bool = False, lrr: bool = False, lrr_monitor: str = 'loss', lrr_factor: float = 0.8, lrr_patience: int = 3, lrr_verbose: int = 1, lrr_mode: str = 'auto', lrr_cooldown: int = 5, lrr_min_learning_rate: float = 1e-06, dense_sizes: list | None = None, dropout_rate: list | None = None, ppm_sizes: list | None = None, activations: list | None = None, repeat_end: list | None = None, dl1_units: list | None = None, dropout_rate_phi: list | None = None, dropout_rate_f: list | None = None, dips_ppm_units: list | None = None, dips_dense_units: list | None = None, dips_loss_weight: float | None = None, intermediate_units: list | None = None, dips_ppm_condition: bool | None = None, ppm_condition: bool | None = None, pooling: str = 'attention', attention_sizes: list | None = None, attention_condition: bool | None = None, n_conditions: int = 0, dense_condition: bool | None = None)#

Bases: object

Dataclass for the nn_structure config options.

activations: list = None#

attention_condition: bool = None#

attention_sizes: list = None#

batch_normalisation: bool = False#

batch_size: int = None#

check_class_labels()#

Check the class_labels that no jet is used twice or a wrongly defined class_labels list is used.

Raises:

ValueError – If the label value was used twice.
ValueError – When the default and extended bjets are defined.
ValueError – When the default and extended cjets are defined.

check_n_conditions()#

Check that the n_conditions are correctly defined together with the tagger.

Raises:

ValueError – If n_conditions is == 0 and the tagger is cads.
ValueError – If n_conditions is != 0 and the tagger is dips_attention.

check_options()#

Check the given options for wrongly/mis-defined values/types. Also check that all needed variables are set.

Raises:: ValueError – If the given tagger is not supported.

class_labels: list = None#

dense_condition: bool = None#

dense_sizes: list = None#

dips_dense_units: list = None#

dips_loss_weight: float = None#

dips_ppm_condition: bool = None#

dips_ppm_units: list = None#

dl1_units: list = None#

dropout_rate: list = None#

dropout_rate_f: list = None#

dropout_rate_phi: list = None#

epochs: int = None#

evaluate_trained_model: bool = True#

intermediate_units: list = None#

learning_rate: float = None#

load_optimiser: bool = False#

lrr: bool = False#

lrr_cooldown: int = 5#

lrr_factor: float = 0.8#

lrr_min_learning_rate: float = 1e-06#

lrr_mode: str = 'auto'#

lrr_monitor: str = 'loss'#

lrr_patience: int = 3#

lrr_verbose: int = 1#

main_class: str = None#

n_conditions: int = 0#

n_jets_train: int = None#

nfiles_tfrecord: int = 5#

pooling: str = 'attention'#

ppm_condition: bool = None#

ppm_sizes: list = None#

repeat_end: list = None#

tagger: str = None#

use_sample_weights: bool = False#

class umami.train_tools.configuration.TrainConfiguration(yaml_config: str | None = None, **kwargs)#

Bases: object

docstring for Configuration.

get_configuration()#: Initalise the dataclasses and check the values.

load_config_file()#: Load config file from disk.

class umami.train_tools.configuration.TrainConfigurationObject(model_name: str | None = None, preprocess_config: str | None = None, evaluate_trained_model: bool = True, tracks_name: str | None = None, train_file: str | None = None, train_data_structure: str | None = None, validation_files: dict | None = None, test_files: dict | None = None, model_file: str | None = None, continue_training: bool = False, exclude: list | None = None)#

Bases: object

Dataclass for the global train config options.

continue_training: bool = False#

evaluate_trained_model: bool = True#

exclude: list = None#

model_file: str = None#

model_name: str = None#

preprocess_config: str = None#

test_files: dict = None#

tracks_name: str = None#

train_data_structure: str = None#

train_file: str = None#

validation_files: dict = None#

class umami.train_tools.configuration.ValidationSettingsConfig(n_jets: int | None = None, working_point: float | None = None, plot_datatype: str = 'pdf', taggers_from_file: dict | None = None, tagger_label: str | None = None, trained_taggers: dict | None = None, use_atlas_tag: bool = True, atlas_first_tag: str = 'Simulation Internal', atlas_second_tag: str = '$\\sqrt{s}=13$ TeV, PFlow jets', val_batch_size: int | None = None, figsize: list = <factory>)#

Bases: object

Dataclass for the validation_settings config options.

atlas_first_tag: str = 'Simulation Internal'#

atlas_second_tag: str = '$\\sqrt{s}=13$ TeV, PFlow jets'#

figsize: list#

n_jets: int = None#

property plot_args#: Plotting arguments for the training plots :returns: Arguments for plotting API :rtype: dict

plot_datatype: str = 'pdf'#

tagger_label: str = None#

taggers_from_file: dict = None#

trained_taggers: dict = None#

use_atlas_tag: bool = True#

val_batch_size: int = None#

working_point: float = None#

umami.train_tools.nn_tools module#

Helper functions for training tools.

class umami.train_tools.nn_tools.CallbackBase(class_labels: list, main_class: str, model_name: str, use_lrr: bool, n_jets: int | None = None, val_data_dict: dict | None = None, target_beff: float = 0.77, frac_dict: dict | None = None, continue_training: bool = False, batch_size: int = 15000)#

Bases: Callback

Base class for the callbacks of the different models.

This class provides the base functionalites for the different callbacks needed for the models that are available.

class umami.train_tools.nn_tools.MyCallback(class_labels: list, main_class: str, model_name: str, use_lrr: bool, n_jets: int | None = None, val_data_dict: dict | None = None, target_beff: float = 0.77, frac_dict: dict | None = None, continue_training: bool = False, batch_size: int = 15000)#

Bases: CallbackBase

Callback class for the standard taggers

This class is the callback for the standard taggers. Only one output (not like the umami tagger) is given.

on_epoch_end(epoch: int, logs: dict | None = None) → None#

Get the needed metrics at epoch end and calculate rest.

This method saves the training metrics at the end of the epoch and also calculates the validation metrics and the rejections for each non-main class for given efficiency and fraction values. Those are also saved.

Parameters:

epoch (int) – Number of the epoch which just finished and is now evaluated and saved.
logs (dict) – Dict with the training metrics of the just finished epoch.

class umami.train_tools.nn_tools.MyCallbackUmami(class_labels: list, main_class: str, model_name: str, use_lrr: bool, n_jets: int | None = None, val_data_dict: dict | None = None, target_beff: float = 0.77, frac_dict: dict | None = None, continue_training: bool = False, batch_size: int = 15000)#

Bases: CallbackBase

Callback class for the umami tagger

This class is the callback for the umami tagger. Due to the two outputs of the tagger, we need special metrics etc.

on_epoch_end(epoch: int, logs: dict | None = None) → None#

Get the needed metrics at epoch end and calculate rest.

Parameters:

epoch (int) – Number of the epoch which just finished and is now evaluated and saved.
logs (dict) – Dict with the training metrics of the just finished epoch.

umami.train_tools.nn_tools.calc_validation_metrics(train_config: object, tagger: str, target_beff: float = 0.77, n_jets: int = 300000, model_string: str = 'model_epoch') → str#

Calculates the validation metrics and rejections for each epoch and dump it into a json.

Parameters:

train_config (object) – The loaded train config object.
tagger (str) – Name of the tagger that is used to calcualte metrics.
target_beff (float) – Working point that is to be used.
n_jets (int) – Number of jets to use for calculation.
model_string (str) – Name of the model files.

Returns:

Path to the validation dict where the results are saved in.

Return type:

output_file_path

Raises:

ValueError – If “tagger” is not dips, dl1, umami or cads.

umami.train_tools.nn_tools.create_metadata_folder(train_config_path: str, var_dict_path: str, model_name: str, preprocess_config_path: str, model_file_path: str | None = None, overwrite_config: bool = False) → None#

Create a metadata folder in the new model_name dir and copy all configs there and change the paths inside the configs to the new metadata directory path. The following configs are copied: preprocessing config, variable dict, scale dict, training config and model_file

Parameters:

train_config_path (str) – Path to the train config that is used.
var_dict_path (str) – Path to the variable dict that is used.
model_name (str) – Model name that is used.
preprocess_config_path (str) – Path to the preprocessing config that is used.
model_file_path (str) – Path to a model to start from (the model given in model_file).
overwrite_config (bool) – If configs already in metadata folder, overwrite them or not.

umami.train_tools.nn_tools.evaluate_model(model: object, data_dict: dict, class_labels: list, main_class: str, target_beff: float = 0.77, frac_dict: dict | None = None, batch_size: int = 15000) → dict#

Evaluate the DIPS/DL1 model on the data provided.

Parameters:

model (object) – Loaded UMAMI model for evaluation.
data_dict (dict) – Dict with the loaded data which are to be evaluated.
class_labels (list) – List of classes used for training of the model.
main_class (str) – Main class which is to be tagged.
target_beff (float) – Working Point which is to be used for evaluation.
frac_dict (dict) – Dict with the fractions of the non-main classes. Sum needs to be one!
batch_size (int, optional) – Number of jets used for validation. By default 15_000

Returns:

result_dict – Dict with validation metrics/rejections.

Return type:

dict

umami.train_tools.nn_tools.evaluate_model_umami(model: object, data_dict: dict, class_labels: list, main_class: str, frac_dict: dict, target_beff: float = 0.77, batch_size: int = 15000) → dict#

Evaluate the UMAMI model on the data provided.

Parameters:

model (object) – Loaded UMAMI model for evaluation.
data_dict (dict) – Dict with the loaded data which are to be evaluated.
class_labels (list) – List of classes used for training of the model.
main_class (str) – Main class which is to be tagged.
frac_dict (dict) – Dict with the fractions of the non-main classes. Sum needs to be one!
target_beff (float, optional) – Working Point which is to be used for evaluation. By default 0.77
batch_size (int, optional) – Number of jets used for validation. By default 15_000

Returns:

result_dict – Dict with validation metrics/rejections.

Return type:

dict

umami.train_tools.nn_tools.get_dropout_rates(dropout_rates_name: str, layer_sizes_name: str, config: dict)#

Helper function to obtain the dropout rates from the training config (with a check if it has same size as the array of layer sizes of the neural network).

Parameters:

dropout_rates_name (str) – Name of the key for the dropout rates (in the training config)
layer_sizes_name (str) – Name of the key for the layer sizes (in the training config)
config (dict) – Dict that stores these keys

Returns:

List with the dropout rates for the layers of the neural network

Return type:

list

Raises:

ValueError – If the number of specified dropout rates != number of layers

umami.train_tools.nn_tools.get_epoch_from_string(string)#

Get the epoch from the model file string.

Parameters:: string (str) – Path of the model file.
Returns:: epoch – Epoch of the model file.
Return type:: int

umami.train_tools.nn_tools.get_jet_feature_indices(variable_header: dict, exclude: list | None = None)#

Deletes from the jet samples the keys listed in exclude.

Parameters:

variable_header (dict) – List with the variables.
exclude (list) – List with the variables that are to be excluded.

Returns:

variables (list) – List with the new variables without the excluded ones.
excluded_variables (list) – List of the excluded variables.
excluded_var_indices (list) – List of the indicies of the excluded variables.

umami.train_tools.nn_tools.get_jet_feature_position(variable_list: list, column_names: list) → list#

Return the index position of the variables listed in variable_list within

the column_names list. WARNING: should match the column order of the training data!

Parameters:

variable_list (list) – List with the variables
column_names (list) – List with the names of the columns

Returns:

List with the positions of the columns

Return type:

list

Raises:

ValueError – If the variable is not in the set.

umami.train_tools.nn_tools.get_metrics_file_name(working_point: float, n_jets: int, dir_name: str)#

Get the validation dict name based on WP, number of jets and dir_name.

Parameters:

working_point (float) – Working point that was used to calculate validation dict.
n_jets (int) – Number of jets that was used to calculate validation dict.
dir_name (str) – Path of the directory where the validation dict is saved.

Returns:

train_dict_path (str) – Path of the training dict.
validation_dict_path (str) – Path of the validation dict.

umami.train_tools.nn_tools.get_model_path(model_name: str, epoch: int) → str#

Get the path where the model will be saved/is saved.

Parameters:

model_name (str) – Name of the model that is to be saved/loaded.
epoch (int) – The epoch which is to be saved/loaded

Returns:

model_path – Path to the model file of the specified epoch.

Return type:

str

umami.train_tools.nn_tools.get_parameters_from_validation_dict_name(dict_name: str) → dict#

Get the parameters used to calculate the validation dict from the validation dict name.

Parameters:: dict_name (str) – Name of the validation dict.
Returns:: parameters – Dict with the parameters (WP, n_jets, dir_name) used to calculate the validation dict.
Return type:: dict
Raises:: Exception – If the name of the validation dict could be rebuild from the extracted parameters.

umami.train_tools.nn_tools.get_test_file(input_file: str, var_dict: str, scale_dict: str, class_labels: list, tracks_name: str, n_jets: int | None = None, indices_to_load: tuple | None = None, exclude: list | None = None, cut_vars_dict: dict | None = None, jet_variables: list | None = None, print_logger: bool = True)#

Load the jet and track variables and labels. Scale the jet and track variables for validation use in the NN’s.

Parameters:

input_file (str) – Path to the file which is to be loaded.
var_dict (str) – Path to the variable dict or the loaded variable dict directly.
scale_dict (str) – Path to the scale dict or the loaded scale dict directly.
class_labels (list) – List of classes used for training of the model.
tracks_name (str) – Name of the tracks collection to use.
n_jets (int, optional) – Number of jets that should be loaded. By default None
indices_to_load (tuple, optional) – Tuple of the indices which are to be loaded. First value is the lower index, second value is the upper index. By default None.
exclude (list) – List of variables that are not loaded.
cut_vars_dict (dict) – Dict with the cuts that should be applied.
jet_variables (list) – List of variables that are used.
print_logger (bool) – Decide, if the logger info is printed or not.

Returns:

x_test (numpy.ndarray) – X values of the jets ready to be used in the NN’s.
x_trk (numpy.ndarray) – X values of the tracks ready to be used in the NN’s.
y_test (numpy.ndarray) – Y values ready to be used in the NN’s.

umami.train_tools.nn_tools.get_test_sample(input_file: str, var_dict: str, scale_dict: str, class_labels: list, n_jets: int | None = None, indices_to_load: tuple | None = None, exclude: list | None = None, cut_vars_dict: dict | None = None, jet_variables: list | None = None, print_logger: bool = True)#

Load the jet variables and labels. Scale the jet variables for validation use in the NN’s.

Parameters:

input_file (str) – Path to the file which is to be loaded.
var_dict (str) – Path to the variable dict or the loaded variable dict directly.
scale_dict (str) – Path to the scale dict or the loaded scale dict directly.
class_labels (list) – List of classes used for training of the model.
n_jets (int, optional) – Number of jets that should be loaded. By default None
indices_to_load (tuple, optional) – Tuple of the indices which are to be loaded. First value is the lower index, second value is the upper index. By default None.
exclude (list) – List of variables that are not loaded.
cut_vars_dict (dict) – Dict with the cuts that should be applied.
jet_variables (list) – List of variables that are used.
print_logger (bool) – Decide, if the logger info is printed or not.

Returns:

jets (numpy.ndarray) – X values of the jets ready to be used in the NN’s.
labels (numpy.ndarray) – Y values ready to be used in the NN’s.

Raises:

ValueError – If jet_variables and exclude are used at the same time.
RuntimeError – If no input file could be found in the given filepath.

umami.train_tools.nn_tools.get_test_sample_trks(input_file: str, var_dict: str, scale_dict: str, class_labels: list, tracks_name: str, n_jets: int | None = None, indices_to_load: tuple | None = None, cut_vars_dict: dict | None = None, print_logger: bool = False)#

Load the track variables and labels. Scale the track variables for validation use in the NN’s.

Parameters:

input_file (str) – Path to the file which is to be loaded.
var_dict (str) – Path to the variable dict or the loaded variable dict directly.
scale_dict (str) – Path to the scale dict or the loaded scale dict directly.
class_labels (list) – List of classes used for training of the model.
tracks_name (str) – Name of tracks collection to use.
n_jets (int, optional) – Number of jets that should be loaded. By default None
indices_to_load (tuple, optional) – Tuple of the indices which are to be loaded. First value is the lower index, second value is the upper index. By default None.
cut_vars_dict (dict) – Dict with the cuts that should be applied.
print_logger (bool) – Decide, if the logger info is printed or not.

Returns:

trks (numpy.ndarray) – X values of the tracks ready to be used in the NN’s.
binary_labels (numpy.ndarray) – Y values ready to be used in the NN’s.

Raises:

RuntimeError – If no input file could be found in the given filepath.

umami.train_tools.nn_tools.get_unique_identifiers(keys: list, prefix: str) → list#

Helper function which takes a list of strings, searches them for a given prefix of the form “prefix_<identifier>” and returns the remaining part of the matching strings

Parameters:

keys (list) – List of strings which are searched for the given prefix
prefix (str) – Prefix to search for in the provided strings

Returns:

identifiers – Sorted list of the unique identifiers that could be found.

Return type:

list

umami.train_tools.nn_tools.load_validation_data(train_config: object, n_jets: int, jets_var_list: list | None = None, convert_to_tensor: bool = False, n_cond: int | None = None) → dict#

Load the validation data for UMAMI.

Parameters:

train_config (object) – Loaded train_config object.
n_jets (int) – Number of jets to load.
jets_var_list (list) – List with jet variables that are to be loaded.
convert_to_tensor (bool) – Decide, if the validation data are converted to tensorflow tensors to avoid memory leaks.
n_cond (int) – Number of addittional variables used for attention

Returns:

val_data_dict – Dict with the validation data.

Return type:

dict

Raises:

ValueError – If the given tagger is not supported.

umami.train_tools.nn_tools.setup_output_directory(dir_name: str, continue_training: bool = False) → None#

Check the output directory path and init/clean it.

Parameters:

dir_name (str) – Path of the output directory.
continue_training (bool) – Decide, if the training is continued (True) and the files in the folder are not erased or a new training is started (False) and the old model files and validation json files are erased.

Raises:

FileExistsError – If the dir_name is an existing file.

umami.train_tools package#

Submodules#

umami.train_tools.configuration module#

umami.train_tools.nn_tools module#

Module contents#