utils.py

#!/usr/bin/env python3
 

from yaml import full_load
from basf2 import find_file
 
 
def get_variables(particle_list, ranked_variable, variables=None, particleNumber=1):
    """ creates variable name pattern requested by the basf2 variable getVariableByRank()
    :param particle_list:
    :param ranked_variable:
    :param variables:
    :param particleNumber:
    :return:
    """
    var_list = []
    for i_num in range(1, particleNumber + 1):
        for var in variables:
            var_list.append(f'getVariableByRank({particle_list}, {ranked_variable}, {var}, {str(i_num)})')
    return var_list
 
 
def load_config(uniqueIdentifier):
    """
    Loads the correct yaml config file for the specified TFlaT uniqueIdentifier.
    All dynamic alterations of the config are handeled in this function.
    The final config is returned as a dict.
    """
    with open(find_file(f'analysis/data/{uniqueIdentifier}.yaml')) as f:
        config = full_load(f)
 
    # Calculate input list lengths
    config["parameters"]["num_trk_features"] = len(config["trk_variable_list"])
    config["parameters"]["num_ecl_features"] = len(config["ecl_variable_list"])
    config["parameters"]["num_roe_features"] = len(config["roe_variable_list"])
 
    return config