study_da

GenerateScan

A class to generate a study (along with the corresponding tree) from a parameter file, and potentially a set of template files.

Attributes:

Name	Type	Description
config	dict	The configuration dictionary.
ryaml	YAML	The YAML parser.
dic_common_parameters	dict	Dictionary of common parameters across generations.

Methods:

Name	Description
__init__	Initializes the generation scan with a configuration file or dictionary.
render	Renders the study file using a template.
write	Writes the study file to disk.
generate_render_write	Generates, renders, and writes the study file.
get_dic_parametric_scans	Retrieves dictionaries of parametric scan values.
parse_parameter_space	Parses the parameter space for a given parameter.
browse_and_collect_parameter_space	Browses and collects the parameter space for a given generation.
postprocess_parameter_lists	Postprocesses the parameter lists.
create_scans	Creates study files for parametric scans.
complete_tree	Completes the tree structure of the study dictionary.
write_tree	Writes the study tree structure to a YAML file.
create_study_for_current_gen	Creates study files for the current generation.
create_study	Creates study files for the entire study.
eval_conditions	Evaluates the conditions to filter out some parameter values.
filter_for_concomitant_parameters	Filters the conditions for concomitant parameters.

Source code in study_da/generate/generate_scan.py

class GenerateScan:
    """
    A class to generate a study (along with the corresponding tree) from a parameter file,
    and potentially a set of template files.

    Attributes:
        config (dict): The configuration dictionary.
        ryaml (yaml.YAML): The YAML parser.
        dic_common_parameters (dict): Dictionary of common parameters across generations.

    Methods:
        __init__(): Initializes the generation scan with a configuration file or dictionary.
        render(): Renders the study file using a template.
        write(): Writes the study file to disk.
        generate_render_write(): Generates, renders, and writes the study file.
        get_dic_parametric_scans(): Retrieves dictionaries of parametric scan values.
        parse_parameter_space(): Parses the parameter space for a given parameter.
        browse_and_collect_parameter_space(): Browses and collects the parameter space for a given
            generation.
        postprocess_parameter_lists(): Postprocesses the parameter lists.
        create_scans(): Creates study files for parametric scans.
        complete_tree(): Completes the tree structure of the study dictionary.
        write_tree(): Writes the study tree structure to a YAML file.
        create_study_for_current_gen(): Creates study files for the current generation.
        create_study(): Creates study files for the entire study.
        eval_conditions(): Evaluates the conditions to filter out some parameter values.
        filter_for_concomitant_parameters(): Filters the conditions for concomitant parameters.
    """

    def __init__(
        self, path_config: Optional[str] = None, dic_scan: Optional[dict[str, Any]] = None
    ):  # sourcery skip: remove-redundant-if
        """
        Initialize the generation scan with a configuration file or dictionary.

        Args:
            path_config (Optional[str]): Path to the configuration file for the scan.
                Default is None.
            dic_scan (Optional[dict[str, Any]]): Dictionary containing the scan configuration.
                Default is None.

        Raises:
            ValueError: If neither or both of `path_config` and `dic_scan` are provided.
        """
        # Load the study configuration from file or dictionary
        if dic_scan is None and path_config is None:
            raise ValueError(
                "Either a path to the configuration file or a dictionary must be provided."
            )
        elif dic_scan is not None and path_config is not None:
            raise ValueError("Only one of the configuration file or dictionary must be provided.")
        elif path_config is not None:
            self.config, self.ryaml = load_dic_from_path(path_config)
        elif dic_scan is not None:
            self.config = dic_scan
            self.ryaml = yaml.YAML()
        else:
            raise ValueError("An unexpected error occurred.")

        # Parameters common across all generations (e.g. for parallelization)
        self.dic_common_parameters: dict[str, Any] = {}

        # Path to the tree file
        self.path_tree = self.config["name"] + "/" + "tree.yaml"

    def render(
        self,
        str_parameters: str,
        template_path: str,
        path_main_configuration: str,
        study_path: Optional[str] = None,
        str_dependencies: Optional[dict[str, str]] = None,
    ) -> str:
        """
        Renders the study file using a template.

        Args:
            str_parameters (str): The string representation of parameters to declare/mutate.
            template_path (str): The path to the template file.
            path_main_configuration (str): The path to the main configuration file.
            study_path (str, optional): The path to the root of the study. Defaults to None.
            dependencies (dict[str, str], optional): The dictionary of dependencies. Defaults to {}.

        Returns:
            str: The rendered study file.
        """

        # Handle mutable default argument
        if str_dependencies is None:
            dependencies = ""
        if study_path is None:
            study_path = ""

        # Generate generations from template
        directory_path = os.path.dirname(template_path)
        template_name = os.path.basename(template_path)
        environment = Environment(
            loader=FileSystemLoader(directory_path),
            variable_start_string="{}  ###---",
            variable_end_string="---###",
        )
        template = environment.get_template(template_name)

        # Better not to render the dependencies path this way, as it becomes too cumbersome to
        # handle the paths when using clusters

        return template.render(
            parameters=str_parameters,
            main_configuration=path_main_configuration,
            path_root_study=study_path,
            # dependencies = str_dependencies,
        )

    def write(self, study_str: str, file_path: str, format_with_black: bool = True):
        """
        Writes the study file to disk.

        Args:
            study_str (str): The study file string.
            file_path (str): The path to write the study file.
            format_with_black (bool, optional): Whether to format the output file with black.
                Defaults to True.
        """

        # Format the string with black
        if format_with_black:
            study_str = format_str(study_str, mode=FileMode())

        # Make folder if it doesn't exist
        folder = os.path.dirname(file_path)
        if folder != "":
            os.makedirs(folder, exist_ok=True)

        with open(file_path, mode="w", encoding="utf-8") as file:
            file.write(study_str)

    def generate_render_write(
        self,
        gen_name: str,
        job_directory_path: str,
        template_path: str,
        depth_gen: int,
        dic_mutated_parameters: dict[str, Any] = {},
    ) -> list[str]:  # sourcery skip: default-mutable-arg
        """
        Generates, renders, and writes the study file.

        Args:
            gen_name (str): The name of the generation.
            study_path (str): The path to the job folder.
            template_path (str): The path to the template folder.
            depth_gen (int): The depth of the generation in the tree.
            dic_mutated_parameters (dict[str, Any], optional): The dictionary of mutated parameters.
                Defaults to {}.

        Returns:
            tuple[str, list[str]]: The study file string and the list of study paths.
        """

        directory_path_gen = f"{job_directory_path}"
        if not directory_path_gen.endswith("/"):
            directory_path_gen += "/"
        file_path_gen = f"{directory_path_gen}{gen_name}.py"
        logging.info(f'Now rendering generation "{file_path_gen}"')

        # Generate the string of parameters
        str_parameters = "{"
        for key, value in dic_mutated_parameters.items():
            if isinstance(value, str):
                str_parameters += f"'{key}' : '{value}', "
            else:
                str_parameters += f"'{key}' : {value}, "
        str_parameters += "}"

        # Adapt the dict of dependencies to the current generation
        dic_dependencies = self.config["dependencies"] if "dependencies" in self.config else {}

        # Unpacking list of dependencies
        dic_dependencies = {
            **{
                key: value for key, value in dic_dependencies.items() if not isinstance(value, list)
            },
            **{
                f"{key}_{str(i).zfill(len(str(len(value))))}": i_value
                for key, value in dic_dependencies.items()
                if isinstance(value, list)
                for i, i_value in enumerate(value)
            },
        }
        self.config["dependencies"] = dic_dependencies

        # Initial dependencies are always copied at the root of the study (hence value.split("/")[-1])
        dic_dependencies = {
            key: "../" * depth_gen + value.split("/")[-1] for key, value in dic_dependencies.items()
        }

        # Always load configuration from above generation, and remove the path from dependencies
        path_main_configuration = "../" + dic_dependencies.pop("main_configuration").split("/")[-1]

        # Create the str for the dependencies
        str_dependencies = "{"
        for key, value in dic_dependencies.items():
            str_dependencies += f"'{key}' : '{value}', "
        str_dependencies += "}"

        # Render and write the study file
        study_str = self.render(
            str_parameters,
            template_path=template_path,
            path_main_configuration=path_main_configuration,
            study_path=os.path.abspath(self.config["name"]),
            str_dependencies=str_dependencies,
        )

        self.write(study_str, file_path_gen)
        return [directory_path_gen]

    def get_dic_parametric_scans(
        self, generation: str
    ) -> tuple[dict[str, Any], dict[str, Any], np.ndarray | None]:
        """
        Retrieves dictionaries of parametric scan values.

        Args:
            generation: The generation name.

        Returns:
            tuple[dict[str, Any], dict[str, Any], np.ndarray|None]: The dictionaries of parametric
                scan values, another dictionnary with better naming for the tree creation, and an
                array of conditions to filter out some parameter values.
        """

        if generation == "base":
            raise ValueError("Generation 'base' should not have scans.")

        # Remember common parameters as they might be used across generations
        if "common_parameters" in self.config["structure"][generation]:
            self.dic_common_parameters[generation] = {}
            for parameter in self.config["structure"][generation]["common_parameters"]:
                self.dic_common_parameters[generation][parameter] = self.config["structure"][
                    generation
                ]["common_parameters"][parameter]

        # Check that the generation has scans
        if (
            "scans" not in self.config["structure"][generation]
            or self.config["structure"][generation]["scans"] is None
        ):
            dic_parameter_lists = {"": [generation]}
            dic_parameter_lists_for_naming = {"": [generation]}
            array_conditions = None
            ll_concomitant_parameters = []
        else:
            # Browse and collect the parameter space for the generation
            (
                dic_parameter_lists,
                dic_parameter_lists_for_naming,
                dic_subvariables,
                ll_concomitant_parameters,
                l_conditions,
            ) = self.browse_and_collect_parameter_space(generation)

            # Get the dimension corresponding to each parameter
            dic_dimension_indices = {
                parameter: idx for idx, parameter in enumerate(dic_parameter_lists)
            }

            # Generate array of conditions to filter out some of the values later
            # Is an array of True values if no conditions are present
            array_conditions = self.eval_conditions(l_conditions, dic_parameter_lists)

            # Filter for concomitant parameters
            array_conditions = self.filter_for_concomitant_parameters(
                array_conditions, ll_concomitant_parameters, dic_dimension_indices
            )

            # Postprocess the parameter lists and update the dictionaries
            dic_parameter_lists, dic_parameter_lists_for_naming = self.postprocess_parameter_lists(
                dic_parameter_lists, dic_parameter_lists_for_naming, dic_subvariables
            )

        return (
            dic_parameter_lists,
            dic_parameter_lists_for_naming,
            array_conditions,
        )

    def parse_parameter_space(
        self,
        parameter: str,
        dic_curr_parameter: dict[str, Any],
        dic_parameter_lists: dict[str, Any],
        dic_parameter_lists_for_naming: dict[str, Any],
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """
        Parses the parameter space for a given parameter.

        Args:
            parameter (str): The parameter name.
            dic_curr_parameter (dict[str, Any]): The dictionary of current parameter values.
            dic_parameter_lists (dict[str, Any]): The dictionary of parameter lists.
            dic_parameter_lists_for_naming (dict[str, Any]): The dictionary of parameter lists for naming.

        Returns:
            tuple[dict[str, Any], dict[str, Any]]: The updated dictionaries of parameter lists.
        """

        if "linspace" in dic_curr_parameter:
            parameter_list = linspace(dic_curr_parameter["linspace"])
            dic_parameter_lists_for_naming[parameter] = parameter_list
        elif "logspace" in dic_curr_parameter:
            parameter_list = logspace(dic_curr_parameter["logspace"])
            dic_parameter_lists_for_naming[parameter] = parameter_list
        elif "path_list" in dic_curr_parameter:
            l_values_path_list = dic_curr_parameter["path_list"]
            parameter_list = list_values_path(l_values_path_list, self.dic_common_parameters)
            dic_parameter_lists_for_naming[parameter] = [
                f"{n:02d}" for n, path in enumerate(parameter_list)
            ]
        elif "list" in dic_curr_parameter:
            parameter_list = dic_curr_parameter["list"]
            dic_parameter_lists_for_naming[parameter] = parameter_list
        elif "expression" in dic_curr_parameter:
            parameter_list = np.round(
                eval(dic_curr_parameter["expression"], copy.deepcopy(dic_parameter_lists)),
                8,
            )
            dic_parameter_lists_for_naming[parameter] = parameter_list
        else:
            raise ValueError(f"Scanning method for parameter {parameter} is not recognized.")

        dic_parameter_lists[parameter] = np.array(parameter_list)
        return dic_parameter_lists, dic_parameter_lists_for_naming

    def browse_and_collect_parameter_space(
        self,
        generation: str,
    ) -> tuple[
        dict[str, Any],
        dict[str, Any],
        dict[str, Any],
        list[list[str]],
        list[str],
    ]:
        """
        Browses and collects the parameter space for a given generation.

        Args:
            generation (str): The generation name.

        Returns:
            tuple[dict[str, Any], dict[str, Any], dict[str, Any], list[list[str]]]: The updated
                dictionaries of parameter lists.
        """

        l_conditions = []
        ll_concomitant_parameters = []
        dic_subvariables = {}
        dic_parameter_lists = {}
        dic_parameter_lists_for_naming = {}
        for parameter in self.config["structure"][generation]["scans"]:
            dic_curr_parameter = self.config["structure"][generation]["scans"][parameter]

            # Parse the parameter space
            dic_parameter_lists, dic_parameter_lists_for_naming = self.parse_parameter_space(
                parameter, dic_curr_parameter, dic_parameter_lists, dic_parameter_lists_for_naming
            )

            # Store potential subvariables
            if "subvariables" in dic_curr_parameter:
                dic_subvariables[parameter] = dic_curr_parameter["subvariables"]

            # Save the condition if it exists
            if "condition" in dic_curr_parameter:
                l_conditions.append(dic_curr_parameter["condition"])

            # Save the concomitant parameters if they exist
            if "concomitant" in dic_curr_parameter:
                if not isinstance(dic_curr_parameter["concomitant"], list):
                    dic_curr_parameter["concomitant"] = [dic_curr_parameter["concomitant"]]
                for concomitant_parameter in dic_curr_parameter["concomitant"]:
                    # Assert that the parameters list have the same size
                    assert len(dic_parameter_lists[parameter]) == len(
                        dic_parameter_lists[concomitant_parameter]
                    ), (
                        f"Parameters {parameter} and {concomitant_parameter} must have the "
                        "same size."
                    )
                # Add to the list for filtering later
                ll_concomitant_parameters.append([parameter] + dic_curr_parameter["concomitant"])

        return (
            dic_parameter_lists,
            dic_parameter_lists_for_naming,
            dic_subvariables,
            ll_concomitant_parameters,
            l_conditions,
        )

    def postprocess_parameter_lists(
        self,
        dic_parameter_lists: dict[str, Any],
        dic_parameter_lists_for_naming: dict[str, Any],
        dic_subvariables: dict[str, Any],
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """
        Post-processes parameter lists by ensuring values are not numpy types and handling nested
        parameters.

        Args:
            dic_parameter_lists (dict[str, Any]): Dictionary containing parameter lists.
            dic_parameter_lists_for_naming (dict[str, Any]): Dictionary containing parameter lists
                for naming.
            dic_subvariables (dict[str, Any]): Dictionary containing subvariables for nested
                parameters.

        Returns:
            tuple[dict[str, Any], dict[str, Any]]: Updated dictionaries of parameter lists and
                parameter lists for naming.
        """
        for parameter, parameter_list in dic_parameter_lists.items():
            parameter_list_for_naming = dic_parameter_lists_for_naming[parameter]

            # Ensure that all values are not numpy types (to avoid serialization issues)
            parameter_list = [x.item() if isinstance(x, np.generic) else x for x in parameter_list]

            # Handle nested parameters
            parameter_list_updated = (
                convert_for_subvariables(dic_subvariables[parameter], parameter_list)
                if parameter in dic_subvariables
                else parameter_list
            )
            # Update the dictionaries
            dic_parameter_lists[parameter] = parameter_list_updated
            dic_parameter_lists_for_naming[parameter] = parameter_list_for_naming

        return dic_parameter_lists, dic_parameter_lists_for_naming

    def create_scans(
        self,
        generation: str,
        generation_path: str,
        template_path: str,
        depth_gen: int,
        dic_parameter_lists: Optional[dict[str, Any]] = None,
        dic_parameter_lists_for_naming: Optional[dict[str, Any]] = None,
        add_prefix_to_folder_names: bool = False,
    ) -> list[str]:
        """
        Creates study files for parametric scans.
        If a dictionary of parameter lists is provided, the scan will be done on the parameter
        lists (no cartesian product). Otherwise, the scan will be done on the cartesian product of
        the parameters defined in the scan configuration file.

        Args:
            generation (str): The generation name.
            generation_path (str): The (relative) path to the generation folder.
            template_path (str): The path to the template folder.
            depth_gen (int): The depth of the generation in the tree.
            dic_parameter_lists (Optional[dict[str, Any]]): The dictionary of parameter lists.
                Defaults to None.
            dic_parameter_lists_for_naming (Optional[dict[str, Any]]): The dictionary of parameter
                lists for naming. Defaults to None.
            add_prefix_to_folder_names (bool): Whether to add a prefix to the folder names. Defaults
                to False.

        Returns:
            tuple[list[str], list[str]]: The list of study file strings and the list of study paths.
        """
        if dic_parameter_lists is None:
            # Get dictionnary of parametric values being scanned
            dic_parameter_lists, dic_parameter_lists_for_naming, array_conditions = (
                self.get_dic_parametric_scans(generation)
            )
            apply_cartesian_product = True
        else:
            if dic_parameter_lists_for_naming is None:
                dic_parameter_lists_for_naming = copy.deepcopy(dic_parameter_lists)
            array_conditions = None
            apply_cartesian_product = False

        # Generate render write for the parameters parameters
        l_study_path = []
        if apply_cartesian_product:
            logging.info(
                f"Now generation cartesian product of all parameters for generation: {generation}"
            )
            array_param_values = itertools.product(*dic_parameter_lists.values())
            array_param_values_for_naming = itertools.product(
                *dic_parameter_lists_for_naming.values()
            )
            array_idx = itertools.product(*[range(len(x)) for x in dic_parameter_lists.values()])
        else:
            logging.info(f"Now generation parameters for generation: {generation}")
            array_param_values = [list(x) for x in zip(*dic_parameter_lists.values())]
            array_param_values_for_naming = [
                list(x) for x in zip(*dic_parameter_lists_for_naming.values())
            ]
            array_idx = range(len(array_param_values))

        # Loop over the parameters
        to_disk_len = np.sum(array_conditions) if array_conditions is not None else 1
        to_disk_idx = 0
        for idx, (l_values, l_values_for_naming, l_idx) in enumerate(
            zip(array_param_values, array_param_values_for_naming, array_idx)
        ):
            # Check the idx to keep if conditions are present
            if array_conditions is not None and not array_conditions[l_idx]:
                continue

            # Create the path for the study
            dic_mutated_parameters = dict(zip(dic_parameter_lists.keys(), l_values))
            dic_mutated_parameters_for_naming = dict(
                zip(dic_parameter_lists.keys(), l_values_for_naming)
            )

            # Handle prefix
            prefix_path = ""
            if add_prefix_to_folder_names:
                prefix_path = f"ID_{str(to_disk_idx).zfill(len(str(to_disk_len)))}_"
                to_disk_idx += 1

            # Handle suffix
            suffix_path = "_".join(
                [
                    f"{parameter}_{value}"
                    for parameter, value in dic_mutated_parameters_for_naming.items()
                ]
            )
            suffix_path = suffix_path.removeprefix("_")

            # Create final path
            path = generation_path + prefix_path + suffix_path + "/"

            # Add common parameters
            if generation in self.dic_common_parameters:
                dic_mutated_parameters |= self.dic_common_parameters[generation]

            # Remove "" from mutated parameters, if it's in the dictionary
            # as it's only used when no scan is done
            if "" in dic_mutated_parameters:
                dic_mutated_parameters.pop("")

            # Generate the study for current generation
            self.generate_render_write(
                generation,
                path,
                template_path,
                depth_gen,
                dic_mutated_parameters=dic_mutated_parameters,
            )

            # Append the list of study paths to build the tree later on
            l_study_path.append(path)

        if not l_study_path:
            logging.warning(
                f"No study paths were created for generation {generation}."
                "Please check the conditions."
            )

        return l_study_path

    def complete_tree(
        self, dictionary_tree: dict, l_study_path_next_gen: list[str], gen: str
    ) -> dict:
        """
        Completes the tree structure of the study dictionary.

        Args:
            dictionary_tree (dict): The dictionary representing the study tree structure.
            l_study_path_next_gen (list[str]): The list of study paths for the next gen.
            gen (str): The generation name.

        Returns:
            dict: The updated dictionary representing the study tree structure.
        """
        logging.info(f"Completing the tree structure for generation: {gen}")
        for path_next in l_study_path_next_gen:
            nested_set(
                dictionary_tree,
                path_next.split("/")[1:-1] + [gen],
                {"file": f"{path_next}{gen}.py"},
            )

        return dictionary_tree

    def write_tree(self, dictionary_tree: dict):
        """
        Writes the study tree structure to a YAML file.

        Args:
            dictionary_tree (dict): The dictionary representing the study tree structure.
        """
        logging.info("Writing the tree structure to a YAML file.")
        ryaml = yaml.YAML()
        with open(self.path_tree, "w") as yaml_file:
            ryaml.indent(sequence=4, offset=2)
            ryaml.dump(dictionary_tree, yaml_file)

    def create_study_for_current_gen(
        self,
        generation: str,
        generation_path: str,
        depth_gen: int,
        dic_parameter_lists: Optional[dict[str, Any]] = None,
        dic_parameter_lists_for_naming: Optional[dict[str, Any]] = None,
        add_prefix_to_folder_names: bool = False,
    ) -> list[str]:
        """
        Creates study files for the current generation.

        Args:
            generation (str): The name of the current generation.
            directory_path (str): The (relative) path to the directory folder for the current
                generation.
            depth_gen (int): The depth of the generation in the tree.
            dic_parameter_lists (Optional[dict[str, Any]]): The dictionary of parameter lists.
                Defaults to None.
            dic_parameter_lists_for_naming (Optional[dict[str, Any]]): The dictionary of parameter
                lists for naming. Defaults to None.
            add_prefix_to_folder_names (bool): Whether to add a prefix to the folder names. Defaults
                to False.

        Returns:
            tuple[list[str], list[str]]: The list of study file strings and the list of study paths.
        """
        executable_path = self.config["structure"][generation]["executable"]
        path_local_template = (
            f"{os.path.dirname(inspect.getfile(GenerateScan))}/../assets/template_scripts/"
        )

        # Check if the executable path corresponds to a file
        if not os.path.isfile(executable_path):
            # Check if the executable path corresponds to a file in the template folder
            executable_path_template = f"{path_local_template}{executable_path}"
            if not os.path.isfile(executable_path_template):
                raise FileNotFoundError(
                    f"Executable file {executable_path} not found locally nor in the study-da "
                    "template folder."
                )
            else:
                executable_path = executable_path_template

        # Ensure that the values in dic_parameter_lists can be dumped with ryaml
        if dic_parameter_lists is not None:
            # Recursively convert all numpy types to standard types
            clean_dic(dic_parameter_lists)
            logging.info("An external dictionary of parameters was provided.")
        else:
            logging.info("Creating the dictionnary of parameters from the configuration file.")

        return self.create_scans(
            generation,
            generation_path,
            executable_path,
            depth_gen,
            dic_parameter_lists,
            dic_parameter_lists_for_naming,
            add_prefix_to_folder_names,
        )

    def browse_and_creat_study(
        self,
        dic_parameter_all_gen: Optional[dict[str, dict[str, Any]]],
        dic_parameter_all_gen_naming: Optional[dict[str, dict[str, Any]]],
        add_prefix_to_folder_names: bool,
    ) -> dict:
        l_study_path = [self.config["name"] + "/"]
        dictionary_tree = {}

        # Browse through the generations
        l_generations = list(self.config["structure"].keys())
        for idx, generation in enumerate(l_generations):
            l_study_path_all_next_generation = []
            logging.info(f"Taking care of generation: {generation}")
            for study_path in l_study_path:
                if dic_parameter_all_gen is None or generation not in dic_parameter_all_gen:
                    dic_parameter_current_gen = None
                    dic_parameter_naming_current_gen = None
                else:
                    dic_parameter_current_gen = dic_parameter_all_gen[generation]
                    if (
                        dic_parameter_all_gen_naming is not None
                        and generation in dic_parameter_all_gen_naming
                    ):
                        dic_parameter_naming_current_gen = dic_parameter_all_gen_naming[generation]
                    else:
                        dic_parameter_naming_current_gen = None

                # Get list of paths for the children of the current study
                l_study_path_next_generation = self.create_study_for_current_gen(
                    generation,
                    study_path,
                    idx + 1,
                    dic_parameter_current_gen,
                    dic_parameter_naming_current_gen,
                    add_prefix_to_folder_names,
                )
                # Update tree
                dictionary_tree = self.complete_tree(
                    dictionary_tree, l_study_path_next_generation, generation
                )
                # Complete list of paths for the children of all studies (of the current generation)
                l_study_path_all_next_generation.extend(l_study_path_next_generation)

            # Update study path for next later
            l_study_path = l_study_path_all_next_generation

        return dictionary_tree

    def create_study(
        self,
        tree_file: bool = True,
        force_overwrite: bool = False,
        dic_parameter_all_gen: Optional[dict[str, dict[str, Any]]] = None,
        dic_parameter_all_gen_naming: Optional[dict[str, dict[str, Any]]] = None,
        add_prefix_to_folder_names: bool = False,
    ) -> None:
        """
        Creates study files for the entire study.

        Args:
            tree_file (bool, optional): Whether to write the study tree structure to a YAML file.
                Defaults to True.
            force_overwrite (bool, optional): Whether to overwrite existing study files.
                Defaults to False.
            dic_parameter_all_gen (Optional[dict[str, dict[str, Any]]]): The dictionary of parameter
                lists for all generations. Defaults to None.
            dic_parameter_all_gen_naming (Optional[dict[str, dict[str, Any]]]): The dictionary of
                parameter lists for all generations for naming. Defaults to None.
            add_prefix_to_folder_names (bool): Whether to add a prefix to the folder names. Defaults
                to False.

        Returns:
            list[str]: The list of study file strings.
        """

        # Raise an error if dic_parameter_all_gen_naming is not None while dic_parameter_all_gen is None
        if dic_parameter_all_gen is None and dic_parameter_all_gen_naming is not None:
            raise ValueError(
                "If dic_parameter_all_gen_naming is defined, dic_parameter_all_gen must be defined."
            )

        # Remove existing study if force_overwrite
        if os.path.exists(self.config["name"]):
            if not force_overwrite:
                logging.info(
                    f"Study {self.config['name']} already exists. Set force_overwrite to True to "
                    "overwrite. Continuing without overwriting."
                )
                return
            shutil.rmtree(self.config["name"])

        # Browse through the generations and create the study
        dictionary_tree = self.browse_and_creat_study(
            dic_parameter_all_gen,
            dic_parameter_all_gen_naming,
            add_prefix_to_folder_names,
        )

        # Add dependencies to root of the study
        if "dependencies" in self.config:
            for dependency, path in self.config["dependencies"].items():
                # Check if the dependency exists as a file
                if not os.path.isfile(path):
                    # Check if the dependency exists as a file in the template folder
                    path_template = f"{os.path.dirname(inspect.getfile(GenerateScan))}/../assets/configurations/{path}"
                    if not os.path.isfile(path_template):
                        raise FileNotFoundError(
                            f"Dependency file {path} not found locally nor in the study-da "
                            "template folder."
                        )
                    else:
                        path = path_template
                shutil.copy2(path, self.config["name"])

        if tree_file:
            self.write_tree(dictionary_tree)

    @staticmethod
    def eval_conditions(l_condition: list[str], dic_parameter_lists: dict[str, Any]) -> np.ndarray:
        """
        Evaluates the conditions to filter out some parameter values.

        Args:
            l_condition (list[str]): The list of conditions.
            dic_parameter_lists (dict[str: Any]): The dictionary of parameter lists.

        Returns:
            np.ndarray: The array of conditions.
        """
        # Initialize the array of parameters as a meshgrid of all parameters
        l_parameters = list(dic_parameter_lists.values())
        meshgrid = np.meshgrid(*l_parameters, indexing="ij")

        # Associate the parameters to their names
        dic_param_mesh = dict(zip(dic_parameter_lists.keys(), meshgrid))

        # Evaluate the conditions and take the intersection of all conditions
        array_conditions = np.ones_like(meshgrid[0], dtype=bool)
        for condition in l_condition:
            array_conditions = array_conditions & eval(condition, dic_param_mesh)

        return array_conditions

    @staticmethod
    def filter_for_concomitant_parameters(
        array_conditions: np.ndarray,
        ll_concomitant_parameters: list[list[str]],
        dic_dimension_indices: dict[str, int],
    ) -> np.ndarray:
        """
        Filters the conditions for concomitant parameters.

        Args:
            array_conditions (np.ndarray): The array of conditions.
            ll_concomitant_parameters (list[list[str]]): The list of concomitant parameters.
            dic_dimension_indices (dict[str, int]): The dictionary of dimension indices.

        Returns:
            np.ndarray: The filtered array of conditions.
        """

        # Return the array of conditions if no concomitant parameters
        if not ll_concomitant_parameters:
            return array_conditions

        # Get the indices of the concomitant parameters
        ll_idx_concomitant_parameters = [
            [dic_dimension_indices[parameter] for parameter in concomitant_parameters]
            for concomitant_parameters in ll_concomitant_parameters
        ]

        # Browse all the values of array_conditions
        for idx, _ in np.ndenumerate(array_conditions):
            # Check if the value is on the diagonal of the concomitant parameters
            for l_idx_concomitant_parameter in ll_idx_concomitant_parameters:
                if any(
                    idx[i] != idx[j]
                    for i, j in itertools.combinations(l_idx_concomitant_parameter, 2)
                ):
                    array_conditions[idx] = False
                    break

        return array_conditions

__init__(path_config=None, dic_scan=None)

Initialize the generation scan with a configuration file or dictionary.

Parameters:

Name	Type	Description	Default
path_config	Optional[str]	Path to the configuration file for the scan. Default is None.	None
dic_scan	Optional[dict[str, Any]]	Dictionary containing the scan configuration. Default is None.	None

Raises:

Type	Description
ValueError	If neither or both of path_config and dic_scan are provided.

Source code in study_da/generate/generate_scan.py

def __init__(
    self, path_config: Optional[str] = None, dic_scan: Optional[dict[str, Any]] = None
):  # sourcery skip: remove-redundant-if
    """
    Initialize the generation scan with a configuration file or dictionary.

    Args:
        path_config (Optional[str]): Path to the configuration file for the scan.
            Default is None.
        dic_scan (Optional[dict[str, Any]]): Dictionary containing the scan configuration.
            Default is None.

    Raises:
        ValueError: If neither or both of `path_config` and `dic_scan` are provided.
    """
    # Load the study configuration from file or dictionary
    if dic_scan is None and path_config is None:
        raise ValueError(
            "Either a path to the configuration file or a dictionary must be provided."
        )
    elif dic_scan is not None and path_config is not None:
        raise ValueError("Only one of the configuration file or dictionary must be provided.")
    elif path_config is not None:
        self.config, self.ryaml = load_dic_from_path(path_config)
    elif dic_scan is not None:
        self.config = dic_scan
        self.ryaml = yaml.YAML()
    else:
        raise ValueError("An unexpected error occurred.")

    # Parameters common across all generations (e.g. for parallelization)
    self.dic_common_parameters: dict[str, Any] = {}

    # Path to the tree file
    self.path_tree = self.config["name"] + "/" + "tree.yaml"

browse_and_collect_parameter_space(generation)

Browses and collects the parameter space for a given generation.

Parameters:

Name	Type	Description	Default
generation	str	The generation name.	required

Returns:

Type	Description
tuple[dict[str, Any], dict[str, Any], dict[str, Any], list[list[str]], list[str]]	tuple[dict[str, Any], dict[str, Any], dict[str, Any], list[list[str]]]: The updated dictionaries of parameter lists.

Source code in study_da/generate/generate_scan.py

def browse_and_collect_parameter_space(
    self,
    generation: str,
) -> tuple[
    dict[str, Any],
    dict[str, Any],
    dict[str, Any],
    list[list[str]],
    list[str],
]:
    """
    Browses and collects the parameter space for a given generation.

    Args:
        generation (str): The generation name.

    Returns:
        tuple[dict[str, Any], dict[str, Any], dict[str, Any], list[list[str]]]: The updated
            dictionaries of parameter lists.
    """

    l_conditions = []
    ll_concomitant_parameters = []
    dic_subvariables = {}
    dic_parameter_lists = {}
    dic_parameter_lists_for_naming = {}
    for parameter in self.config["structure"][generation]["scans"]:
        dic_curr_parameter = self.config["structure"][generation]["scans"][parameter]

        # Parse the parameter space
        dic_parameter_lists, dic_parameter_lists_for_naming = self.parse_parameter_space(
            parameter, dic_curr_parameter, dic_parameter_lists, dic_parameter_lists_for_naming
        )

        # Store potential subvariables
        if "subvariables" in dic_curr_parameter:
            dic_subvariables[parameter] = dic_curr_parameter["subvariables"]

        # Save the condition if it exists
        if "condition" in dic_curr_parameter:
            l_conditions.append(dic_curr_parameter["condition"])

        # Save the concomitant parameters if they exist
        if "concomitant" in dic_curr_parameter:
            if not isinstance(dic_curr_parameter["concomitant"], list):
                dic_curr_parameter["concomitant"] = [dic_curr_parameter["concomitant"]]
            for concomitant_parameter in dic_curr_parameter["concomitant"]:
                # Assert that the parameters list have the same size
                assert len(dic_parameter_lists[parameter]) == len(
                    dic_parameter_lists[concomitant_parameter]
                ), (
                    f"Parameters {parameter} and {concomitant_parameter} must have the "
                    "same size."
                )
            # Add to the list for filtering later
            ll_concomitant_parameters.append([parameter] + dic_curr_parameter["concomitant"])

    return (
        dic_parameter_lists,
        dic_parameter_lists_for_naming,
        dic_subvariables,
        ll_concomitant_parameters,
        l_conditions,
    )

complete_tree(dictionary_tree, l_study_path_next_gen, gen)

Completes the tree structure of the study dictionary.

Parameters:

Name	Type	Description	Default
dictionary_tree	dict	The dictionary representing the study tree structure.	required
l_study_path_next_gen	list[str]	The list of study paths for the next gen.	required
gen	str	The generation name.	required

Returns:

Name	Type	Description
dict	dict	The updated dictionary representing the study tree structure.

Source code in study_da/generate/generate_scan.py

def complete_tree(
    self, dictionary_tree: dict, l_study_path_next_gen: list[str], gen: str
) -> dict:
    """
    Completes the tree structure of the study dictionary.

    Args:
        dictionary_tree (dict): The dictionary representing the study tree structure.
        l_study_path_next_gen (list[str]): The list of study paths for the next gen.
        gen (str): The generation name.

    Returns:
        dict: The updated dictionary representing the study tree structure.
    """
    logging.info(f"Completing the tree structure for generation: {gen}")
    for path_next in l_study_path_next_gen:
        nested_set(
            dictionary_tree,
            path_next.split("/")[1:-1] + [gen],
            {"file": f"{path_next}{gen}.py"},
        )

    return dictionary_tree

create_scans(generation, generation_path, template_path, depth_gen, dic_parameter_lists=None, dic_parameter_lists_for_naming=None, add_prefix_to_folder_names=False)

Creates study files for parametric scans. If a dictionary of parameter lists is provided, the scan will be done on the parameter lists (no cartesian product). Otherwise, the scan will be done on the cartesian product of the parameters defined in the scan configuration file.

Parameters:

Name	Type	Description	Default
generation	str	The generation name.	required
generation_path	str	The (relative) path to the generation folder.	required
template_path	str	The path to the template folder.	required
depth_gen	int	The depth of the generation in the tree.	required
dic_parameter_lists	Optional[dict[str, Any]]	The dictionary of parameter lists. Defaults to None.	None
dic_parameter_lists_for_naming	Optional[dict[str, Any]]	The dictionary of parameter lists for naming. Defaults to None.	None
add_prefix_to_folder_names	bool	Whether to add a prefix to the folder names. Defaults to False.	False

Returns:

Type	Description
list[str]	tuple[list[str], list[str]]: The list of study file strings and the list of study paths.

Source code in study_da/generate/generate_scan.py

def create_scans(
    self,
    generation: str,
    generation_path: str,
    template_path: str,
    depth_gen: int,
    dic_parameter_lists: Optional[dict[str, Any]] = None,
    dic_parameter_lists_for_naming: Optional[dict[str, Any]] = None,
    add_prefix_to_folder_names: bool = False,
) -> list[str]:
    """
    Creates study files for parametric scans.
    If a dictionary of parameter lists is provided, the scan will be done on the parameter
    lists (no cartesian product). Otherwise, the scan will be done on the cartesian product of
    the parameters defined in the scan configuration file.

    Args:
        generation (str): The generation name.
        generation_path (str): The (relative) path to the generation folder.
        template_path (str): The path to the template folder.
        depth_gen (int): The depth of the generation in the tree.
        dic_parameter_lists (Optional[dict[str, Any]]): The dictionary of parameter lists.
            Defaults to None.
        dic_parameter_lists_for_naming (Optional[dict[str, Any]]): The dictionary of parameter
            lists for naming. Defaults to None.
        add_prefix_to_folder_names (bool): Whether to add a prefix to the folder names. Defaults
            to False.

    Returns:
        tuple[list[str], list[str]]: The list of study file strings and the list of study paths.
    """
    if dic_parameter_lists is None:
        # Get dictionnary of parametric values being scanned
        dic_parameter_lists, dic_parameter_lists_for_naming, array_conditions = (
            self.get_dic_parametric_scans(generation)
        )
        apply_cartesian_product = True
    else:
        if dic_parameter_lists_for_naming is None:
            dic_parameter_lists_for_naming = copy.deepcopy(dic_parameter_lists)
        array_conditions = None
        apply_cartesian_product = False

    # Generate render write for the parameters parameters
    l_study_path = []
    if apply_cartesian_product:
        logging.info(
            f"Now generation cartesian product of all parameters for generation: {generation}"
        )
        array_param_values = itertools.product(*dic_parameter_lists.values())
        array_param_values_for_naming = itertools.product(
            *dic_parameter_lists_for_naming.values()
        )
        array_idx = itertools.product(*[range(len(x)) for x in dic_parameter_lists.values()])
    else:
        logging.info(f"Now generation parameters for generation: {generation}")
        array_param_values = [list(x) for x in zip(*dic_parameter_lists.values())]
        array_param_values_for_naming = [
            list(x) for x in zip(*dic_parameter_lists_for_naming.values())
        ]
        array_idx = range(len(array_param_values))

    # Loop over the parameters
    to_disk_len = np.sum(array_conditions) if array_conditions is not None else 1
    to_disk_idx = 0
    for idx, (l_values, l_values_for_naming, l_idx) in enumerate(
        zip(array_param_values, array_param_values_for_naming, array_idx)
    ):
        # Check the idx to keep if conditions are present
        if array_conditions is not None and not array_conditions[l_idx]:
            continue

        # Create the path for the study
        dic_mutated_parameters = dict(zip(dic_parameter_lists.keys(), l_values))
        dic_mutated_parameters_for_naming = dict(
            zip(dic_parameter_lists.keys(), l_values_for_naming)
        )

        # Handle prefix
        prefix_path = ""
        if add_prefix_to_folder_names:
            prefix_path = f"ID_{str(to_disk_idx).zfill(len(str(to_disk_len)))}_"
            to_disk_idx += 1

        # Handle suffix
        suffix_path = "_".join(
            [
                f"{parameter}_{value}"
                for parameter, value in dic_mutated_parameters_for_naming.items()
            ]
        )
        suffix_path = suffix_path.removeprefix("_")

        # Create final path
        path = generation_path + prefix_path + suffix_path + "/"

        # Add common parameters
        if generation in self.dic_common_parameters:
            dic_mutated_parameters |= self.dic_common_parameters[generation]

        # Remove "" from mutated parameters, if it's in the dictionary
        # as it's only used when no scan is done
        if "" in dic_mutated_parameters:
            dic_mutated_parameters.pop("")

        # Generate the study for current generation
        self.generate_render_write(
            generation,
            path,
            template_path,
            depth_gen,
            dic_mutated_parameters=dic_mutated_parameters,
        )

        # Append the list of study paths to build the tree later on
        l_study_path.append(path)

    if not l_study_path:
        logging.warning(
            f"No study paths were created for generation {generation}."
            "Please check the conditions."
        )

    return l_study_path

create_study(tree_file=True, force_overwrite=False, dic_parameter_all_gen=None, dic_parameter_all_gen_naming=None, add_prefix_to_folder_names=False)

Creates study files for the entire study.

Parameters:

Name	Type	Description	Default
tree_file	bool	Whether to write the study tree structure to a YAML file. Defaults to True.	True
force_overwrite	bool	Whether to overwrite existing study files. Defaults to False.	False
dic_parameter_all_gen	Optional[dict[str, dict[str, Any]]]	The dictionary of parameter lists for all generations. Defaults to None.	None
dic_parameter_all_gen_naming	Optional[dict[str, dict[str, Any]]]	The dictionary of parameter lists for all generations for naming. Defaults to None.	None
add_prefix_to_folder_names	bool	Whether to add a prefix to the folder names. Defaults to False.	False

Returns:

Type	Description
None	list[str]: The list of study file strings.

Source code in study_da/generate/generate_scan.py

def create_study(
    self,
    tree_file: bool = True,
    force_overwrite: bool = False,
    dic_parameter_all_gen: Optional[dict[str, dict[str, Any]]] = None,
    dic_parameter_all_gen_naming: Optional[dict[str, dict[str, Any]]] = None,
    add_prefix_to_folder_names: bool = False,
) -> None:
    """
    Creates study files for the entire study.

    Args:
        tree_file (bool, optional): Whether to write the study tree structure to a YAML file.
            Defaults to True.
        force_overwrite (bool, optional): Whether to overwrite existing study files.
            Defaults to False.
        dic_parameter_all_gen (Optional[dict[str, dict[str, Any]]]): The dictionary of parameter
            lists for all generations. Defaults to None.
        dic_parameter_all_gen_naming (Optional[dict[str, dict[str, Any]]]): The dictionary of
            parameter lists for all generations for naming. Defaults to None.
        add_prefix_to_folder_names (bool): Whether to add a prefix to the folder names. Defaults
            to False.

    Returns:
        list[str]: The list of study file strings.
    """

    # Raise an error if dic_parameter_all_gen_naming is not None while dic_parameter_all_gen is None
    if dic_parameter_all_gen is None and dic_parameter_all_gen_naming is not None:
        raise ValueError(
            "If dic_parameter_all_gen_naming is defined, dic_parameter_all_gen must be defined."
        )

    # Remove existing study if force_overwrite
    if os.path.exists(self.config["name"]):
        if not force_overwrite:
            logging.info(
                f"Study {self.config['name']} already exists. Set force_overwrite to True to "
                "overwrite. Continuing without overwriting."
            )
            return
        shutil.rmtree(self.config["name"])

    # Browse through the generations and create the study
    dictionary_tree = self.browse_and_creat_study(
        dic_parameter_all_gen,
        dic_parameter_all_gen_naming,
        add_prefix_to_folder_names,
    )

    # Add dependencies to root of the study
    if "dependencies" in self.config:
        for dependency, path in self.config["dependencies"].items():
            # Check if the dependency exists as a file
            if not os.path.isfile(path):
                # Check if the dependency exists as a file in the template folder
                path_template = f"{os.path.dirname(inspect.getfile(GenerateScan))}/../assets/configurations/{path}"
                if not os.path.isfile(path_template):
                    raise FileNotFoundError(
                        f"Dependency file {path} not found locally nor in the study-da "
                        "template folder."
                    )
                else:
                    path = path_template
            shutil.copy2(path, self.config["name"])

    if tree_file:
        self.write_tree(dictionary_tree)

create_study_for_current_gen(generation, generation_path, depth_gen, dic_parameter_lists=None, dic_parameter_lists_for_naming=None, add_prefix_to_folder_names=False)

Creates study files for the current generation.

Parameters:

Name	Type	Description	Default
generation	str	The name of the current generation.	required
directory_path	str	The (relative) path to the directory folder for the current generation.	required
depth_gen	int	The depth of the generation in the tree.	required
dic_parameter_lists	Optional[dict[str, Any]]	The dictionary of parameter lists. Defaults to None.	None
dic_parameter_lists_for_naming	Optional[dict[str, Any]]	The dictionary of parameter lists for naming. Defaults to None.	None
add_prefix_to_folder_names	bool	Whether to add a prefix to the folder names. Defaults to False.	False

Returns:

Type	Description
list[str]	tuple[list[str], list[str]]: The list of study file strings and the list of study paths.

Source code in study_da/generate/generate_scan.py

def create_study_for_current_gen(
    self,
    generation: str,
    generation_path: str,
    depth_gen: int,
    dic_parameter_lists: Optional[dict[str, Any]] = None,
    dic_parameter_lists_for_naming: Optional[dict[str, Any]] = None,
    add_prefix_to_folder_names: bool = False,
) -> list[str]:
    """
    Creates study files for the current generation.

    Args:
        generation (str): The name of the current generation.
        directory_path (str): The (relative) path to the directory folder for the current
            generation.
        depth_gen (int): The depth of the generation in the tree.
        dic_parameter_lists (Optional[dict[str, Any]]): The dictionary of parameter lists.
            Defaults to None.
        dic_parameter_lists_for_naming (Optional[dict[str, Any]]): The dictionary of parameter
            lists for naming. Defaults to None.
        add_prefix_to_folder_names (bool): Whether to add a prefix to the folder names. Defaults
            to False.

    Returns:
        tuple[list[str], list[str]]: The list of study file strings and the list of study paths.
    """
    executable_path = self.config["structure"][generation]["executable"]
    path_local_template = (
        f"{os.path.dirname(inspect.getfile(GenerateScan))}/../assets/template_scripts/"
    )

    # Check if the executable path corresponds to a file
    if not os.path.isfile(executable_path):
        # Check if the executable path corresponds to a file in the template folder
        executable_path_template = f"{path_local_template}{executable_path}"
        if not os.path.isfile(executable_path_template):
            raise FileNotFoundError(
                f"Executable file {executable_path} not found locally nor in the study-da "
                "template folder."
            )
        else:
            executable_path = executable_path_template

    # Ensure that the values in dic_parameter_lists can be dumped with ryaml
    if dic_parameter_lists is not None:
        # Recursively convert all numpy types to standard types
        clean_dic(dic_parameter_lists)
        logging.info("An external dictionary of parameters was provided.")
    else:
        logging.info("Creating the dictionnary of parameters from the configuration file.")

    return self.create_scans(
        generation,
        generation_path,
        executable_path,
        depth_gen,
        dic_parameter_lists,
        dic_parameter_lists_for_naming,
        add_prefix_to_folder_names,
    )

eval_conditions(l_condition, dic_parameter_lists) staticmethod

Evaluates the conditions to filter out some parameter values.

Parameters:

Name	Type	Description	Default
l_condition	list[str]	The list of conditions.	required
dic_parameter_lists	dict[str	Any]): The dictionary of parameter lists.	required

Returns:

Type	Description
ndarray	np.ndarray: The array of conditions.

Source code in study_da/generate/generate_scan.py

@staticmethod
def eval_conditions(l_condition: list[str], dic_parameter_lists: dict[str, Any]) -> np.ndarray:
    """
    Evaluates the conditions to filter out some parameter values.

    Args:
        l_condition (list[str]): The list of conditions.
        dic_parameter_lists (dict[str: Any]): The dictionary of parameter lists.

    Returns:
        np.ndarray: The array of conditions.
    """
    # Initialize the array of parameters as a meshgrid of all parameters
    l_parameters = list(dic_parameter_lists.values())
    meshgrid = np.meshgrid(*l_parameters, indexing="ij")

    # Associate the parameters to their names
    dic_param_mesh = dict(zip(dic_parameter_lists.keys(), meshgrid))

    # Evaluate the conditions and take the intersection of all conditions
    array_conditions = np.ones_like(meshgrid[0], dtype=bool)
    for condition in l_condition:
        array_conditions = array_conditions & eval(condition, dic_param_mesh)

    return array_conditions

filter_for_concomitant_parameters(array_conditions, ll_concomitant_parameters, dic_dimension_indices) staticmethod

Filters the conditions for concomitant parameters.

Parameters:

Name	Type	Description	Default
array_conditions	ndarray	The array of conditions.	required
ll_concomitant_parameters	list[list[str]]	The list of concomitant parameters.	required
dic_dimension_indices	dict[str, int]	The dictionary of dimension indices.	required

Returns:

Type	Description
ndarray	np.ndarray: The filtered array of conditions.

Source code in study_da/generate/generate_scan.py

@staticmethod
def filter_for_concomitant_parameters(
    array_conditions: np.ndarray,
    ll_concomitant_parameters: list[list[str]],
    dic_dimension_indices: dict[str, int],
) -> np.ndarray:
    """
    Filters the conditions for concomitant parameters.

    Args:
        array_conditions (np.ndarray): The array of conditions.
        ll_concomitant_parameters (list[list[str]]): The list of concomitant parameters.
        dic_dimension_indices (dict[str, int]): The dictionary of dimension indices.

    Returns:
        np.ndarray: The filtered array of conditions.
    """

    # Return the array of conditions if no concomitant parameters
    if not ll_concomitant_parameters:
        return array_conditions

    # Get the indices of the concomitant parameters
    ll_idx_concomitant_parameters = [
        [dic_dimension_indices[parameter] for parameter in concomitant_parameters]
        for concomitant_parameters in ll_concomitant_parameters
    ]

    # Browse all the values of array_conditions
    for idx, _ in np.ndenumerate(array_conditions):
        # Check if the value is on the diagonal of the concomitant parameters
        for l_idx_concomitant_parameter in ll_idx_concomitant_parameters:
            if any(
                idx[i] != idx[j]
                for i, j in itertools.combinations(l_idx_concomitant_parameter, 2)
            ):
                array_conditions[idx] = False
                break

    return array_conditions

generate_render_write(gen_name, job_directory_path, template_path, depth_gen, dic_mutated_parameters={})

Generates, renders, and writes the study file.

Parameters:

Name	Type	Description	Default
gen_name	str	The name of the generation.	required
study_path	str	The path to the job folder.	required
template_path	str	The path to the template folder.	required
depth_gen	int	The depth of the generation in the tree.	required
dic_mutated_parameters	dict[str, Any]	The dictionary of mutated parameters. Defaults to {}.	{}

Returns:

Type	Description
list[str]	tuple[str, list[str]]: The study file string and the list of study paths.

Source code in study_da/generate/generate_scan.py

def generate_render_write(
    self,
    gen_name: str,
    job_directory_path: str,
    template_path: str,
    depth_gen: int,
    dic_mutated_parameters: dict[str, Any] = {},
) -> list[str]:  # sourcery skip: default-mutable-arg
    """
    Generates, renders, and writes the study file.

    Args:
        gen_name (str): The name of the generation.
        study_path (str): The path to the job folder.
        template_path (str): The path to the template folder.
        depth_gen (int): The depth of the generation in the tree.
        dic_mutated_parameters (dict[str, Any], optional): The dictionary of mutated parameters.
            Defaults to {}.

    Returns:
        tuple[str, list[str]]: The study file string and the list of study paths.
    """

    directory_path_gen = f"{job_directory_path}"
    if not directory_path_gen.endswith("/"):
        directory_path_gen += "/"
    file_path_gen = f"{directory_path_gen}{gen_name}.py"
    logging.info(f'Now rendering generation "{file_path_gen}"')

    # Generate the string of parameters
    str_parameters = "{"
    for key, value in dic_mutated_parameters.items():
        if isinstance(value, str):
            str_parameters += f"'{key}' : '{value}', "
        else:
            str_parameters += f"'{key}' : {value}, "
    str_parameters += "}"

    # Adapt the dict of dependencies to the current generation
    dic_dependencies = self.config["dependencies"] if "dependencies" in self.config else {}

    # Unpacking list of dependencies
    dic_dependencies = {
        **{
            key: value for key, value in dic_dependencies.items() if not isinstance(value, list)
        },
        **{
            f"{key}_{str(i).zfill(len(str(len(value))))}": i_value
            for key, value in dic_dependencies.items()
            if isinstance(value, list)
            for i, i_value in enumerate(value)
        },
    }
    self.config["dependencies"] = dic_dependencies

    # Initial dependencies are always copied at the root of the study (hence value.split("/")[-1])
    dic_dependencies = {
        key: "../" * depth_gen + value.split("/")[-1] for key, value in dic_dependencies.items()
    }

    # Always load configuration from above generation, and remove the path from dependencies
    path_main_configuration = "../" + dic_dependencies.pop("main_configuration").split("/")[-1]

    # Create the str for the dependencies
    str_dependencies = "{"
    for key, value in dic_dependencies.items():
        str_dependencies += f"'{key}' : '{value}', "
    str_dependencies += "}"

    # Render and write the study file
    study_str = self.render(
        str_parameters,
        template_path=template_path,
        path_main_configuration=path_main_configuration,
        study_path=os.path.abspath(self.config["name"]),
        str_dependencies=str_dependencies,
    )

    self.write(study_str, file_path_gen)
    return [directory_path_gen]

get_dic_parametric_scans(generation)

Retrieves dictionaries of parametric scan values.

Parameters:

Name	Type	Description	Default
generation	str	The generation name.	required

Returns:

Type	Description
tuple[dict[str, Any], dict[str, Any], ndarray | None]	tuple[dict[str, Any], dict[str, Any], np.ndarray|None]: The dictionaries of parametric scan values, another dictionnary with better naming for the tree creation, and an array of conditions to filter out some parameter values.

Source code in study_da/generate/generate_scan.py

def get_dic_parametric_scans(
    self, generation: str
) -> tuple[dict[str, Any], dict[str, Any], np.ndarray | None]:
    """
    Retrieves dictionaries of parametric scan values.

    Args:
        generation: The generation name.

    Returns:
        tuple[dict[str, Any], dict[str, Any], np.ndarray|None]: The dictionaries of parametric
            scan values, another dictionnary with better naming for the tree creation, and an
            array of conditions to filter out some parameter values.
    """

    if generation == "base":
        raise ValueError("Generation 'base' should not have scans.")

    # Remember common parameters as they might be used across generations
    if "common_parameters" in self.config["structure"][generation]:
        self.dic_common_parameters[generation] = {}
        for parameter in self.config["structure"][generation]["common_parameters"]:
            self.dic_common_parameters[generation][parameter] = self.config["structure"][
                generation
            ]["common_parameters"][parameter]

    # Check that the generation has scans
    if (
        "scans" not in self.config["structure"][generation]
        or self.config["structure"][generation]["scans"] is None
    ):
        dic_parameter_lists = {"": [generation]}
        dic_parameter_lists_for_naming = {"": [generation]}
        array_conditions = None
        ll_concomitant_parameters = []
    else:
        # Browse and collect the parameter space for the generation
        (
            dic_parameter_lists,
            dic_parameter_lists_for_naming,
            dic_subvariables,
            ll_concomitant_parameters,
            l_conditions,
        ) = self.browse_and_collect_parameter_space(generation)

        # Get the dimension corresponding to each parameter
        dic_dimension_indices = {
            parameter: idx for idx, parameter in enumerate(dic_parameter_lists)
        }

        # Generate array of conditions to filter out some of the values later
        # Is an array of True values if no conditions are present
        array_conditions = self.eval_conditions(l_conditions, dic_parameter_lists)

        # Filter for concomitant parameters
        array_conditions = self.filter_for_concomitant_parameters(
            array_conditions, ll_concomitant_parameters, dic_dimension_indices
        )

        # Postprocess the parameter lists and update the dictionaries
        dic_parameter_lists, dic_parameter_lists_for_naming = self.postprocess_parameter_lists(
            dic_parameter_lists, dic_parameter_lists_for_naming, dic_subvariables
        )

    return (
        dic_parameter_lists,
        dic_parameter_lists_for_naming,
        array_conditions,
    )

parse_parameter_space(parameter, dic_curr_parameter, dic_parameter_lists, dic_parameter_lists_for_naming)

Parses the parameter space for a given parameter.

Parameters:

Name	Type	Description	Default
parameter	str	The parameter name.	required
dic_curr_parameter	dict[str, Any]	The dictionary of current parameter values.	required
dic_parameter_lists	dict[str, Any]	The dictionary of parameter lists.	required
dic_parameter_lists_for_naming	dict[str, Any]	The dictionary of parameter lists for naming.	required

Returns:

Type	Description
tuple[dict[str, Any], dict[str, Any]]	tuple[dict[str, Any], dict[str, Any]]: The updated dictionaries of parameter lists.

Source code in study_da/generate/generate_scan.py

def parse_parameter_space(
    self,
    parameter: str,
    dic_curr_parameter: dict[str, Any],
    dic_parameter_lists: dict[str, Any],
    dic_parameter_lists_for_naming: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any]]:
    """
    Parses the parameter space for a given parameter.

    Args:
        parameter (str): The parameter name.
        dic_curr_parameter (dict[str, Any]): The dictionary of current parameter values.
        dic_parameter_lists (dict[str, Any]): The dictionary of parameter lists.
        dic_parameter_lists_for_naming (dict[str, Any]): The dictionary of parameter lists for naming.

    Returns:
        tuple[dict[str, Any], dict[str, Any]]: The updated dictionaries of parameter lists.
    """

    if "linspace" in dic_curr_parameter:
        parameter_list = linspace(dic_curr_parameter["linspace"])
        dic_parameter_lists_for_naming[parameter] = parameter_list
    elif "logspace" in dic_curr_parameter:
        parameter_list = logspace(dic_curr_parameter["logspace"])
        dic_parameter_lists_for_naming[parameter] = parameter_list
    elif "path_list" in dic_curr_parameter:
        l_values_path_list = dic_curr_parameter["path_list"]
        parameter_list = list_values_path(l_values_path_list, self.dic_common_parameters)
        dic_parameter_lists_for_naming[parameter] = [
            f"{n:02d}" for n, path in enumerate(parameter_list)
        ]
    elif "list" in dic_curr_parameter:
        parameter_list = dic_curr_parameter["list"]
        dic_parameter_lists_for_naming[parameter] = parameter_list
    elif "expression" in dic_curr_parameter:
        parameter_list = np.round(
            eval(dic_curr_parameter["expression"], copy.deepcopy(dic_parameter_lists)),
            8,
        )
        dic_parameter_lists_for_naming[parameter] = parameter_list
    else:
        raise ValueError(f"Scanning method for parameter {parameter} is not recognized.")

    dic_parameter_lists[parameter] = np.array(parameter_list)
    return dic_parameter_lists, dic_parameter_lists_for_naming

postprocess_parameter_lists(dic_parameter_lists, dic_parameter_lists_for_naming, dic_subvariables)

Post-processes parameter lists by ensuring values are not numpy types and handling nested parameters.

Parameters:

Name	Type	Description	Default
dic_parameter_lists	dict[str, Any]	Dictionary containing parameter lists.	required
dic_parameter_lists_for_naming	dict[str, Any]	Dictionary containing parameter lists for naming.	required
dic_subvariables	dict[str, Any]	Dictionary containing subvariables for nested parameters.	required

Returns:

Type	Description
tuple[dict[str, Any], dict[str, Any]]	tuple[dict[str, Any], dict[str, Any]]: Updated dictionaries of parameter lists and parameter lists for naming.

Source code in study_da/generate/generate_scan.py

def postprocess_parameter_lists(
    self,
    dic_parameter_lists: dict[str, Any],
    dic_parameter_lists_for_naming: dict[str, Any],
    dic_subvariables: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any]]:
    """
    Post-processes parameter lists by ensuring values are not numpy types and handling nested
    parameters.

    Args:
        dic_parameter_lists (dict[str, Any]): Dictionary containing parameter lists.
        dic_parameter_lists_for_naming (dict[str, Any]): Dictionary containing parameter lists
            for naming.
        dic_subvariables (dict[str, Any]): Dictionary containing subvariables for nested
            parameters.

    Returns:
        tuple[dict[str, Any], dict[str, Any]]: Updated dictionaries of parameter lists and
            parameter lists for naming.
    """
    for parameter, parameter_list in dic_parameter_lists.items():
        parameter_list_for_naming = dic_parameter_lists_for_naming[parameter]

        # Ensure that all values are not numpy types (to avoid serialization issues)
        parameter_list = [x.item() if isinstance(x, np.generic) else x for x in parameter_list]

        # Handle nested parameters
        parameter_list_updated = (
            convert_for_subvariables(dic_subvariables[parameter], parameter_list)
            if parameter in dic_subvariables
            else parameter_list
        )
        # Update the dictionaries
        dic_parameter_lists[parameter] = parameter_list_updated
        dic_parameter_lists_for_naming[parameter] = parameter_list_for_naming

    return dic_parameter_lists, dic_parameter_lists_for_naming

render(str_parameters, template_path, path_main_configuration, study_path=None, str_dependencies=None)

Renders the study file using a template.

Parameters:

Name	Type	Description	Default
str_parameters	str	The string representation of parameters to declare/mutate.	required
template_path	str	The path to the template file.	required
path_main_configuration	str	The path to the main configuration file.	required
study_path	str	The path to the root of the study. Defaults to None.	None
dependencies	dict[str, str]	The dictionary of dependencies. Defaults to {}.	required

Returns:

Name	Type	Description
str	str	The rendered study file.

Source code in study_da/generate/generate_scan.py

def render(
    self,
    str_parameters: str,
    template_path: str,
    path_main_configuration: str,
    study_path: Optional[str] = None,
    str_dependencies: Optional[dict[str, str]] = None,
) -> str:
    """
    Renders the study file using a template.

    Args:
        str_parameters (str): The string representation of parameters to declare/mutate.
        template_path (str): The path to the template file.
        path_main_configuration (str): The path to the main configuration file.
        study_path (str, optional): The path to the root of the study. Defaults to None.
        dependencies (dict[str, str], optional): The dictionary of dependencies. Defaults to {}.

    Returns:
        str: The rendered study file.
    """

    # Handle mutable default argument
    if str_dependencies is None:
        dependencies = ""
    if study_path is None:
        study_path = ""

    # Generate generations from template
    directory_path = os.path.dirname(template_path)
    template_name = os.path.basename(template_path)
    environment = Environment(
        loader=FileSystemLoader(directory_path),
        variable_start_string="{}  ###---",
        variable_end_string="---###",
    )
    template = environment.get_template(template_name)

    # Better not to render the dependencies path this way, as it becomes too cumbersome to
    # handle the paths when using clusters

    return template.render(
        parameters=str_parameters,
        main_configuration=path_main_configuration,
        path_root_study=study_path,
        # dependencies = str_dependencies,
    )

write(study_str, file_path, format_with_black=True)

Writes the study file to disk.

Parameters:

Name	Type	Description	Default
study_str	str	The study file string.	required
file_path	str	The path to write the study file.	required
format_with_black	bool	Whether to format the output file with black. Defaults to True.	True

Source code in study_da/generate/generate_scan.py

def write(self, study_str: str, file_path: str, format_with_black: bool = True):
    """
    Writes the study file to disk.

    Args:
        study_str (str): The study file string.
        file_path (str): The path to write the study file.
        format_with_black (bool, optional): Whether to format the output file with black.
            Defaults to True.
    """

    # Format the string with black
    if format_with_black:
        study_str = format_str(study_str, mode=FileMode())

    # Make folder if it doesn't exist
    folder = os.path.dirname(file_path)
    if folder != "":
        os.makedirs(folder, exist_ok=True)

    with open(file_path, mode="w", encoding="utf-8") as file:
        file.write(study_str)

write_tree(dictionary_tree)

Writes the study tree structure to a YAML file.

Parameters:

Name	Type	Description	Default
dictionary_tree	dict	The dictionary representing the study tree structure.	required

Source code in study_da/generate/generate_scan.py

def write_tree(self, dictionary_tree: dict):
    """
    Writes the study tree structure to a YAML file.

    Args:
        dictionary_tree (dict): The dictionary representing the study tree structure.
    """
    logging.info("Writing the tree structure to a YAML file.")
    ryaml = yaml.YAML()
    with open(self.path_tree, "w") as yaml_file:
        ryaml.indent(sequence=4, offset=2)
        ryaml.dump(dictionary_tree, yaml_file)

SubmitScan

Source code in study_da/submit/submit_scan.py

class SubmitScan:
    def __init__(
        self,
        path_tree: str,
        path_python_environment: str = "",
        path_python_environment_container: str = "",
        path_container_image: Optional[str] = None,
    ) -> None:
        """
        Initializes the SubmitScan class.

        Args:
            path_tree (str): The path to the tree structure.
            path_python_environment (str): The path to the Python environment. Defaults to "".
            path_python_environment_container (str, optional): The path to the Python environment
                in the container. Defaults to "".
            path_container_image (Optional[str], optional): The path to the container image.
                Defaults to None.
        """
        # Path to study files
        self.path_tree = path_tree

        # Absolute path to the tree
        self.abs_path_tree = os.path.abspath(path_tree)

        # Name of the study folder
        self.study_name = os.path.dirname(path_tree)

        # Absolute path to the study folder (get from the path_tree)
        self.abs_path = os.path.abspath(self.study_name).split(f"/{self.study_name}")[0]

        # Check that the current working directory is one step above the study folder
        if os.getcwd() != self.abs_path:
            raise ValueError(
                "The current working directory must be the parent folder of the study folder, "
                "i.e. the folder from which the study was generated. "
                "Please submit from there."
            )

        # Container image (Docker or Singularity, if any)
        # Turn to absolute path if it is not already
        if path_container_image is None:
            self.path_container_image = None
        elif not os.path.isabs(path_container_image):
            self.path_container_image = os.path.abspath(path_container_image)
        else:
            self.path_container_image = path_container_image

        # Python environment for the container
        self.path_python_environment_container = path_python_environment_container

        # Ensure that the container image is set if the python environment is set
        if self.path_container_image and not self.path_python_environment_container:
            raise ValueError(
                "The path to the python environment in the container must be set if the container"
                "image is set."
            )

        # Add /bin/activate to the path_python_environment if needed
        if not self.path_python_environment_container.endswith("/bin/activate"):
            # Remove potential / at the end of the path
            if (
                self.path_python_environment_container
                and self.path_python_environment_container[-1] == "/"
            ):
                self.path_python_environment_container = self.path_python_environment_container[:-1]
            self.path_python_environment_container += "/bin/activate"

        # Ensure the path to the python environment is not "" if the container image is not set
        if not self.path_container_image and not path_python_environment:
            raise ValueError(
                "The path to the python environment must be set if the container image is not set."
            )

        # Path to the python environment, activate with `source path_python_environment`
        if not path_python_environment:
            logging.warning("No local python environment provided.")
            self.path_python_environment = ""

        else:
            # Ensure that the path is not of the form path/bin/activate environment_name
            split_path = path_python_environment.split(" ")
            real_path = split_path[0]
            env_name = split_path[1] if len(split_path) > 1 else ""

            # Turn to absolute path if it is not already
            self.path_python_environment = (
                real_path if os.path.isabs(real_path) else os.path.abspath(real_path)
            )

            # Add /bin/activate to the path_python_environment if needed
            if "bin/activate" not in self.path_python_environment:
                # Ensure there's no / at the end of the path
                if self.path_python_environment and self.path_python_environment[-1] == "/":
                    self.path_python_environment = self.path_python_environment[:-1]
                self.path_python_environment += "/bin/activate"

            # Add environment name to the path_python_environment if needed
            if env_name:
                self.path_python_environment += f" {env_name}"
        # Lock file to avoid concurrent access (softlock as several platforms are used)
        self.lock = SoftFileLock(f"{self.path_tree}.lock", timeout=60)

    # dic_tree as a property so that it is reloaded every time it is accessed
    @property
    def dic_tree(self) -> dict:
        """
        Loads the dictionary tree from the path.

        Returns:
            dict: The loaded dictionary tree.
        """
        logging.info(f"Loading tree from {self.path_tree}")
        return load_dic_from_path(self.path_tree)[0]

    # Setter for the dic_tree property
    @dic_tree.setter
    def dic_tree(self, value: dict) -> None:
        """
        Writes the dictionary tree to the path.

        Args:
            value (dict): The dictionary tree to write.
        """
        logging.info(f"Writing tree to {self.path_tree}")
        write_dic_to_path(value, self.path_tree)

    def configure_jobs(
        self,
        force_configure: bool = False,
        dic_config_jobs: Optional[dict[str, dict[str, Any]]] = None,
    ) -> None:
        """
        Configures the jobs by modifying the tree structure and creating the run files for each job.

        Args:
            force_configure (bool, optional): Whether to force reconfiguration. Defaults to False.
            dic_config_jobs (Optional[dict[str, dict[str, Any]]], optional): A dictionary containing
                the configuration of the jobs. Defaults to None.
        """
        # Lock since we are modifying the tree
        logging.info("Acquiring lock to configure jobs")
        with self.lock:
            # Get the tree
            dic_tree = self.dic_tree

            # Ensure jobs have not been configured already
            if ("configured" in dic_tree and dic_tree["configured"]) and not force_configure:
                logging.warning("Jobs have already been configured. Skipping.")
                return

            # Configure the jobs (add generation and job keys, set status to "To finish")
            dic_tree = ConfigJobs(dic_tree,starting_depth=-len(Path(self.path_tree).parts) + 2).find_and_configure_jobs(dic_config_jobs)

            # Add the python environment, container image and absolute path of the study to the tree
            dic_tree["python_environment"] = self.path_python_environment
            dic_tree["container_image"] = self.path_container_image
            dic_tree["absolute_path"] = self.abs_path
            dic_tree["status"] = "to_finish"
            dic_tree["configured"] = True

            # Explicitly set the dic_tree property to force rewrite
            self.dic_tree = dic_tree

        logging.info("Jobs have been configured. Lock released.")

    def get_all_jobs(self) -> dict:
        """
        Retrieves all jobs from the configuration, without modifying the tree.

        Returns:
            dict: A dictionary containing all jobs.
        """
        # Get a copy of the tree as it's safer
        with self.lock:
            dic_tree = self.dic_tree
        return ConfigJobs(dic_tree,starting_depth=-len(Path(self.path_tree).parts) + 2).find_all_jobs()

    def generate_run_files(
        self,
        dic_tree: dict[str, Any],
        l_jobs: list[str],
        dic_additional_commands_per_gen: dict[int, str],
        dic_dependencies_per_gen: dict[int, list[str]],
        dic_copy_back_per_gen: dict[int, dict[str, bool]],
        name_config: str,
    ) -> dict:
        """
        Generates run files for the specified jobs.

        Args:
            dic_tree (dict): The dictionary tree structure.
            l_jobs (list[str]): List of jobs to submit.
            dic_additional_commands_per_gen (dict[int, str], optional): Additional commands per
                generation. Defaults to {}.
            dic_dependencies_per_gen (dict[int, list[str]], optional): Dependencies per generation.
                Only used when doing a HTC submission.
            dic_copy_back_per_gen (Optional[dict[int, dict[str, bool]]], optional): A dictionary
                containing the files to copy back per generation. Accepted keys are "parquet",
                "yaml", "txt", "json", "zip" and "all".
            name_config (str, optional): The name of the configuration file for the study.

        Returns:
            dict: The updated dictionary tree structure.
        """

        logging.info("Generating run files for the jobs to submit")
        # Generate the run files for the jobs to submit
        dic_all_jobs = self.get_all_jobs()
        for job in l_jobs:
            l_keys = dic_all_jobs[job]["l_keys"]
            job_name = os.path.basename(job)
            relative_job_folder = os.path.dirname(job)
            absolute_job_folder = f"{self.abs_path}/{relative_job_folder}"
            generation_number = dic_all_jobs[job]["gen"]
            submission_type = nested_get(dic_tree, l_keys + ["submission_type"])
            singularity = "docker" in submission_type
            path_python_environment = (
                self.path_python_environment_container
                if singularity
                else self.path_python_environment
            )

            # Ensure that the run file does not already exist
            if "path_run" in nested_get(dic_tree, l_keys):
                path_run_curr = nested_get(dic_tree, l_keys + ["path_run"])
                if path_run_curr is not None and os.path.exists(path_run_curr):
                    logging.info(f"Run file already exists for job {job}. Skipping.")
                    continue

            # Build l_dependencies and add to the kwargs
            l_dependencies = dic_dependencies_per_gen.get(generation_number, [])

            # Get arguments of current generation
            dic_args = dic_copy_back_per_gen.get(generation_number, {})

            # Mutate the keys
            dic_args = {f"copy_back_{key}": value for key, value in dic_args.items()}

            # Build kwargs for the run file
            kwargs_htc = {
                "l_dependencies": l_dependencies,
                "name_config": name_config,
            } | dic_args

            run_str = generate_run_file(
                absolute_job_folder,
                job_name,
                path_python_environment,
                htc="htc" in submission_type,
                additionnal_command=dic_additional_commands_per_gen.get(generation_number, ""),
                **kwargs_htc,
            )
            # Write the run file
            path_run_job = f"{absolute_job_folder}/run.sh"
            with open(path_run_job, "w") as f:
                f.write(run_str)

            # Change permissions to make the file executable
            os.chmod(path_run_job, 0o755)

            # Record the path to the run file in the tree
            nested_set(dic_tree, l_keys + ["path_run"], path_run_job)

        return dic_tree

    def check_and_update_all_jobs_status(self) -> tuple[dict[str, Any], str]:
        """
        Checks the status of all jobs and updates their status in the job dictionary.

        This method iterates through all jobs, checks if a ".finished" or a ".failed" file exists in
        the job's folder, and updates the job's status accordingly. If at least one job is not
        finished or failed, the overall status is set to "to_finish". If all jobs are finished or
        failed, the overall status is set to "finished".

        Returns:
            tuple[dict[str, Any], str]: A tuple containing:
            - A dictionary with all jobs and their updated statuses.
            - A string representing the final status ("to_finish" or "finished").
        """
        dic_all_jobs = self.get_all_jobs()
        at_least_one_job_to_finish = False
        final_status = "to_finish"
        with self.lock:
            # Get dic tree once to avoid reloading it for every job
            dic_tree = self.dic_tree

            # First pass to update the state of the tree
            for job in dic_all_jobs:
                # Skip jobs that are already finished, failed or unsubmittable
                if nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) in [
                    "finished",
                    "failed",
                    "unsubmittable",
                ]:
                    continue

                # Check the state of the others
                relative_job_folder = os.path.dirname(job)
                absolute_job_folder = f"{self.abs_path}/{relative_job_folder}"
                if os.path.exists(f"{absolute_job_folder}/.finished"):
                    nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "finished")
                # Check if the job failed otherwise (not to resubmit it again)
                elif os.path.exists(f"{absolute_job_folder}/.failed"):
                    nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "failed")
                # else:
                #     at_least_one_job_to_finish = True

            # Second pass to update the state of the tree with unreachable jobs
            dependency_graph = DependencyGraph(dic_tree, dic_all_jobs)
            for job in dic_all_jobs:
                # Get all failed dependencies across the tree
                l_dep_failed = dependency_graph.get_failed_dependency(job)
                if len(l_dep_failed) > 0:
                    nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "unsubmittable")
                elif nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) == "to_submit":
                    at_least_one_job_to_finish = True

            if not at_least_one_job_to_finish:
                # No more jobs to submit so finished
                dic_tree["status"] = final_status = "finished"
                # Last pass to check if all jobs are properly finished
                for job in dic_all_jobs:
                    if nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) != "finished":
                        dic_tree["status"] = final_status = "finished with issues"
                        break

            # Update dic_tree from cluster_submission
            self.dic_tree = dic_tree

        return dic_all_jobs, final_status

    def reset_failed_jobs(self, dic_tree: dict[str, Any]) -> dict[str, Any]:
        """
        Resets the status of jobs that have failed to "to_submit".

        Args:
            dic_tree (dict[str, Any]): The dictionary tree structure.

        Returns:
            dict[str, Any]: The updated dictionary tree structure.
        """

        dic_all_jobs = self.get_all_jobs()
        # First pass to update the state of the tree
        for job in dic_all_jobs:
            # Skip jobs that are not failed
            if nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) != "failed":
                continue

            # Reset the state of the others
            relative_job_folder = os.path.dirname(job)
            absolute_job_folder = f"{self.abs_path}/{relative_job_folder}"

            # Remove failed tag
            if os.path.exists(f"{absolute_job_folder}/.failed"):
                os.remove(f"{absolute_job_folder}/.failed")
            else:
                logging.warning(f"Failed file not found for job {job}.")

            # Remove run file
            if "path_run" in nested_get(dic_tree, dic_all_jobs[job]["l_keys"]):
                path_run_curr = nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["path_run"])
                if path_run_curr is not None and os.path.exists(path_run_curr):
                    os.remove(path_run_curr)
                else:
                    logging.warning(f"Run file not found for job {job}.")

            # Reset the status of the job
            nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "to_submit")

        return dic_tree

    def submit(
        self,
        one_generation_at_a_time: bool = False,
        dic_additional_commands_per_gen: Optional[dict[int, str]] = None,
        dic_dependencies_per_gen: Optional[dict[int, list[str]]] = None,
        dic_copy_back_per_gen: Optional[dict[int, dict[str, bool]]] = None,
        name_config: str = "config.yaml",
        force_submit: bool = False,
    ) -> str:
        """
        Submits the jobs to the cluster. Note that copying back large files (e.g. json colliders)
        can trigger a throttling mechanism in AFS.

        The following arguments are only used for HTC jobs submission:
        - dic_additional_commands_per_gen
        - dic_dependencies_per_gen
        - dic_copy_back_per_gen
        - name_config

        Args:
            one_generation_at_a_time (bool, optional): Whether to submit one full generation at a
                time. Defaults to False.
            dic_additional_commands_per_gen (dict[int, str], optional): Additional commands per
                generation. Defaults to None.
            dic_dependencies_per_gen (dict[int, list[str]], optional): Dependencies per generation.
                Only used when doing a HTC submission. Defaults to None.
            dic_copy_back_per_gen (Optional[dict[int, dict[str, bool]]], optional): A dictionary
                containing the files to copy back per generation. Accepted keys are "parquet",
                "yaml", "txt", "json", "zip" and "all". Defaults to None, corresponding to copying
                back only "light" files, i.e. parquet, yaml and txt.
            name_config (str, optional): The name of the configuration file for the study.
                Defaults to "config.yaml".
            force_submit (bool, optional): If True, jobs are resubmitted even though they failed.
                Defaults to False.

        Returns:
            str: The final status of the jobs.
        """
        # Handle mutable default arguments
        if dic_additional_commands_per_gen is None:
            dic_additional_commands_per_gen = {}
        if dic_dependencies_per_gen is None:
            dic_dependencies_per_gen = {}
        if dic_copy_back_per_gen is None:
            dic_copy_back_per_gen = {}

        # Handle force submit
        if force_submit:
            logging.warning("Forcing resubmission of all failed jobs.")
            with self.lock:
                # Acquire tree from disk
                dic_tree = self.dic_tree

                # Reset the tree by deleting the failed tags
                dic_tree = self.reset_failed_jobs(dic_tree)
                dic_tree["status"] = "to_finish"
                # Write the tree back to disk
                self.dic_tree = dic_tree

        # Update the status of all jobs before submitting
        dic_all_jobs, final_status = self.check_and_update_all_jobs_status()
        if final_status == "finished":
            print("All jobs are finished.")
            return final_status
        elif final_status == "finished with issues":
            print("All jobs are finished but some did not run properly.")
            return final_status

        logging.info("Acquiring lock to submit jobs")
        with self.lock:
            # Get dic tree once to avoid reloading it for every job
            dic_tree = self.dic_tree

            # Submit the jobs
            self._submit(
                dic_tree,
                dic_all_jobs,
                one_generation_at_a_time,
                dic_additional_commands_per_gen,
                dic_dependencies_per_gen,
                dic_copy_back_per_gen,
                name_config,
            )

            # Update dic_tree from cluster_submission
            self.dic_tree = dic_tree
        logging.info("Jobs have been submitted. Lock released.")
        return final_status

    def _submit(
        self,
        dic_tree: dict[str, Any],
        dic_all_jobs: dict[str, dict[str, Any]],
        one_generation_at_a_time: bool,
        dic_additional_commands_per_gen: dict[int, str],
        dic_dependencies_per_gen: dict[int, list[str]],
        dic_copy_back_per_gen: dict[int, dict[str, bool]],
        name_config: str,
    ) -> None:
        """
        Submits the jobs to the cluster.

        Args:
            dic_tree (dict[str, Any]): The dictionary tree structure.
            dic_all_jobs (dict[str, dict[str,Any]]): A dictionary containing all jobs.
            one_generation_at_a_time (bool): Whether to submit one full generation at a time.
            dic_additional_commands_per_gen (dict[int, str], optional): Additional commands per
                generation.

            The following arguments are only used for HTC jobs submission:

            dic_dependencies_per_gen (dict[int, list[str]], optional): Dependencies per generation.
                Only used when doing a HTC submission. Defaults to None.
            dic_copy_back_per_gen (Optional[dict[int, dict[str, bool]]], optional): A dictionary
                containing the files to copy back per generation.
            name_config (str, optional): The name of the configuration file for the study.
        """
        # Collect dict of list of unfinished jobs for every tree branch and every gen
        dic_to_submit_by_gen = {}
        dic_summary_by_gen = {}
        dependency_graph = DependencyGraph(dic_tree, dic_all_jobs)
        for job in dic_all_jobs:
            dic_to_submit_by_gen, dic_summary_by_gen = self._check_job_submit_status(
                job,
                dic_tree,
                dic_all_jobs,
                dic_to_submit_by_gen,
                dic_summary_by_gen,
                dependency_graph,
            )

        # Only keep the topmost generation if one_generation_at_a_time is True
        if one_generation_at_a_time:
            logging.info(
                "Cropping list of jobs to submit to ensure only one generation is submitted at "
                "a time."
            )
            min_gen = min(k for k, l_jobs in dic_to_submit_by_gen.items() if l_jobs)
            dic_to_submit_by_gen = {min_gen: dic_to_submit_by_gen[min_gen]}

        # Convert dic_to_submit_by_gen to contain all requested information
        l_jobs_to_submit = [job for dic_gen in dic_to_submit_by_gen.values() for job in dic_gen]

        # Generate run files for the jobs to submit
        # ! Run files are generated at submit and not at configuration as the configuration
        # ! files are created at the end of each generation
        dic_tree = self.generate_run_files(
            dic_tree,
            l_jobs_to_submit,
            dic_additional_commands_per_gen,
            dic_dependencies_per_gen=dic_dependencies_per_gen,
            dic_copy_back_per_gen=dic_copy_back_per_gen,
            name_config=name_config,
        )

        # Create the ClusterSubmission object
        path_submission_file = f"{self.abs_path}/{self.study_name}/submission/submission_file.sub"
        cluster_submission = ClusterSubmission(
            self.study_name,
            l_jobs_to_submit,
            dic_all_jobs,
            dic_tree,
            path_submission_file,
            self.abs_path,
        )

        # Write and submit the submission files
        logging.info("Writing and submitting submission files")
        dic_submission_files = cluster_submission.write_sub_files(dic_summary_by_gen)

        # Log the state of the jobs
        self.log_jobs_state(dic_summary_by_gen)
        for submission_type, (
            list_of_jobs,
            l_submission_filenames,
        ) in dic_submission_files.items():
            cluster_submission.submit(list_of_jobs, l_submission_filenames, submission_type)

    @staticmethod
    def log_jobs_state(dic_summary_by_gen: dict[int, dict[str, int]]) -> None:
        """
        Logs the state of jobs for each generation.

        Args:
            dic_summary_by_gen (dict): A dictionary where the keys are generation numbers
                and the values are dictionaries summarizing job states.
                Each summary dictionary should contain the following keys:
                - 'to_submit_later': int, number of jobs left to submit later
                - 'running_or_queuing': int, number of jobs running or queuing
                - 'submitted_now': int, number of jobs submitted now
                - 'finished': int, number of jobs finished
                - 'failed': int, number of jobs failed
                - 'dependency_failed': int, number of jobs on hold due to failed dependencies

        Returns:
            None
        """
        print("State of the jobs:")
        for gen, dic_summary in dic_summary_by_gen.items():
            print("********************************")
            print(f"Generation {gen}")
            print(f"Jobs left to submit later: {dic_summary['to_submit_later']}")
            print(f"Jobs running or queuing: {dic_summary['running_or_queuing']}")
            print(f"Jobs submitted now: {dic_summary['submitted_now']}")
            print(f"Jobs finished: {dic_summary['finished']}")
            print(f"Jobs failed: {dic_summary['failed']}")
            print(f"Jobs on hold due to failed dependencies: {dic_summary['dependency_failed']}")
            print("********************************")

    @staticmethod
    def _check_job_submit_status(
        job: str,
        dic_tree: dict[str, Any],
        dic_all_jobs: dict[str, dict[str, Any]],
        dic_to_submit_by_gen: dict[int, list[str]],
        dic_summary_by_gen: dict[int, dict[str, int]],
        dependency_graph: DependencyGraph,
    ) -> tuple[dict[int, list[str]], dict[int, dict[str, int]]]:
        """
        Checks the status and dependencies of a job and updates the submission and summary
        dictionaries.

        Args:
            job (str): The job identifier.
            dic_tree (dict[str, Any]): The dictionary tree structure.
            dic_all_jobs (dict[str, dict[str,Any]]): A dictionary containing all jobs.
            dic_to_submit_by_gen (dict[int, list[str]]): A dictionary where keys are generation
                numbers and values are lists of jobs to submit for each generation.
            dic_summary_by_gen (dict[int, dict[str, int]]): A dictionary where keys are generation
                numbers and values are dictionaries summarizing job states.
            dependency_graph (DependencyGraph): An object to check job dependencies.

        Returns:
            tuple[dict[int, list[str]], dict[int, dict[str, int]]]: Updated dictionaries for jobs to
                submit and job summaries.
        """
        gen = dic_all_jobs[job]["gen"]
        if gen not in dic_to_submit_by_gen:
            dic_to_submit_by_gen[gen] = []
            dic_summary_by_gen[gen] = {
                "finished": 0,
                "failed": 0,
                "dependency_failed": 0,
                "running_or_queuing": 0,
                "submitted_now": 0,
                "to_submit_later": 0,
            }
        logging.info(f"Checking job {job} dependencies and status in tree")
        l_dep = dependency_graph.get_unfinished_dependency(job)
        l_dep_failed = dependency_graph.get_failed_dependency(job)

        # Job will be on hold as it has failed dependencies
        if len(l_dep_failed) > 0:
            logging.warning(
                f"Job {job} has failed dependencies: {l_dep_failed}, it won't be submitted."
            )
            dic_summary_by_gen[gen]["dependency_failed"] += 1

        # Jobs is waiting for dependencies to finish
        elif len(l_dep) > 0:
            dic_summary_by_gen[gen]["to_submit_later"] += 1

        # Job dependencies are ok
        elif len(l_dep) == 0:
            # But job has failed already
            if nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) == "failed":
                dic_summary_by_gen[gen]["failed"] += 1

            # Or job has finished already
            elif nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) == "finished":
                dic_summary_by_gen[gen]["finished"] += 1

            # Else everything is ok, added to the submit dict
            else:
                logging.info(f"Job {job} is added for submission.")
                dic_to_submit_by_gen[gen].append(job)
                # We'll determine which jobs actually have to be submitted and which jobs
                # are running at the end of the function, after querying the cluster or the local pc

        return dic_to_submit_by_gen, dic_summary_by_gen

    def keep_submit_until_done(
        self,
        one_generation_at_a_time: bool = False,
        wait_time: float = 30,
        max_try=100,
        dic_additional_commands_per_gen: Optional[dict[int, str]] = None,
        dic_dependencies_per_gen: Optional[dict[int, list[str]]] = None,
        dic_copy_back_per_gen: Optional[dict[int, dict[str, bool]]] = None,
        name_config: str = "config.yaml",
        force_submit: bool = False,
    ) -> None:
        """
        Keeps submitting jobs until all jobs are finished or failed.

        The following arguments are only used for HTC jobs submission:
        - dic_additional_commands_per_gen
        - dic_dependencies_per_gen
        - dic_copy_back_per_gen
        - name_config

        Args:
            one_generation_at_a_time (bool, optional): Whether to submit one full generation at a
                time. Defaults to False.
            wait_time (float, optional): The wait time between submissions in minutes.
                Defaults to 30.
            max_try (int, optional): The maximum number of tries before stopping the submission.
            dic_additional_commands_per_gen (dict[int, str], optional): Additional commands per
                generation. Defaults to None.
            dic_dependencies_per_gen (dict[int, list[str]], optional): Dependencies per generation.
                Only used when doing a HTC submission. Defaults to None.
            dic_copy_back_per_gen (Optional[dict[int, dict[str, bool]]], optional): A dictionary
                containing the files to copy back per generation. Accepted keys are "parquet",
                "yaml", "txt", "json", "zip" and "all". Defaults to None, corresponding to copying
                back only "light" files, i.e. parquet, yaml and txt.
            name_config (str, optional): The name of the configuration file for the study.
                Defaults to "config.yaml".
            force_submit (bool, optional): If True, jobs are resubmitted even though they failed.
                Defaults to False.


        Returns:
            None
        """
        # Handle mutable default arguments
        if dic_additional_commands_per_gen is None:
            dic_additional_commands_per_gen = {}
        if dic_dependencies_per_gen is None:
            dic_dependencies_per_gen = {}

        if wait_time < 1 / 20:
            logging.warning("Wait time should be at least 10 seconds to prevent locking errors.")
            logging.warning("Setting wait time to 10 seconds.")
            wait_time = 10 / 60

        # I don't need to lock the tree here since the status cheking is read only and
        # the lock is acquired in the submit method for the submission
        while (
            self.submit(
                one_generation_at_a_time,
                dic_additional_commands_per_gen,
                dic_dependencies_per_gen,
                dic_copy_back_per_gen,
                name_config,
                force_submit=force_submit,
            )
            not in ["finished", "finished with issues"]
            and max_try > 0
        ):
            # Wait for a certain amount of time before checking again
            logging.info(f"Waiting {wait_time} minutes before checking again.")
            time.sleep(wait_time * 60)
            max_try -= 1

        if max_try == 0:
            print("Maximum number of tries reached. Stopping submission.")

dic_tree: dict property writable

Loads the dictionary tree from the path.

Returns:

Name	Type	Description
dict	dict	The loaded dictionary tree.

__init__(path_tree, path_python_environment='', path_python_environment_container='', path_container_image=None)

Initializes the SubmitScan class.

Parameters:

Name	Type	Description	Default
path_tree	str	The path to the tree structure.	required
path_python_environment	str	The path to the Python environment. Defaults to "".	''
path_python_environment_container	str	The path to the Python environment in the container. Defaults to "".	''
path_container_image	Optional[str]	The path to the container image. Defaults to None.	None

Source code in study_da/submit/submit_scan.py

def __init__(
    self,
    path_tree: str,
    path_python_environment: str = "",
    path_python_environment_container: str = "",
    path_container_image: Optional[str] = None,
) -> None:
    """
    Initializes the SubmitScan class.

    Args:
        path_tree (str): The path to the tree structure.
        path_python_environment (str): The path to the Python environment. Defaults to "".
        path_python_environment_container (str, optional): The path to the Python environment
            in the container. Defaults to "".
        path_container_image (Optional[str], optional): The path to the container image.
            Defaults to None.
    """
    # Path to study files
    self.path_tree = path_tree

    # Absolute path to the tree
    self.abs_path_tree = os.path.abspath(path_tree)

    # Name of the study folder
    self.study_name = os.path.dirname(path_tree)

    # Absolute path to the study folder (get from the path_tree)
    self.abs_path = os.path.abspath(self.study_name).split(f"/{self.study_name}")[0]

    # Check that the current working directory is one step above the study folder
    if os.getcwd() != self.abs_path:
        raise ValueError(
            "The current working directory must be the parent folder of the study folder, "
            "i.e. the folder from which the study was generated. "
            "Please submit from there."
        )

    # Container image (Docker or Singularity, if any)
    # Turn to absolute path if it is not already
    if path_container_image is None:
        self.path_container_image = None
    elif not os.path.isabs(path_container_image):
        self.path_container_image = os.path.abspath(path_container_image)
    else:
        self.path_container_image = path_container_image

    # Python environment for the container
    self.path_python_environment_container = path_python_environment_container

    # Ensure that the container image is set if the python environment is set
    if self.path_container_image and not self.path_python_environment_container:
        raise ValueError(
            "The path to the python environment in the container must be set if the container"
            "image is set."
        )

    # Add /bin/activate to the path_python_environment if needed
    if not self.path_python_environment_container.endswith("/bin/activate"):
        # Remove potential / at the end of the path
        if (
            self.path_python_environment_container
            and self.path_python_environment_container[-1] == "/"
        ):
            self.path_python_environment_container = self.path_python_environment_container[:-1]
        self.path_python_environment_container += "/bin/activate"

    # Ensure the path to the python environment is not "" if the container image is not set
    if not self.path_container_image and not path_python_environment:
        raise ValueError(
            "The path to the python environment must be set if the container image is not set."
        )

    # Path to the python environment, activate with `source path_python_environment`
    if not path_python_environment:
        logging.warning("No local python environment provided.")
        self.path_python_environment = ""

    else:
        # Ensure that the path is not of the form path/bin/activate environment_name
        split_path = path_python_environment.split(" ")
        real_path = split_path[0]
        env_name = split_path[1] if len(split_path) > 1 else ""

        # Turn to absolute path if it is not already
        self.path_python_environment = (
            real_path if os.path.isabs(real_path) else os.path.abspath(real_path)
        )

        # Add /bin/activate to the path_python_environment if needed
        if "bin/activate" not in self.path_python_environment:
            # Ensure there's no / at the end of the path
            if self.path_python_environment and self.path_python_environment[-1] == "/":
                self.path_python_environment = self.path_python_environment[:-1]
            self.path_python_environment += "/bin/activate"

        # Add environment name to the path_python_environment if needed
        if env_name:
            self.path_python_environment += f" {env_name}"
    # Lock file to avoid concurrent access (softlock as several platforms are used)
    self.lock = SoftFileLock(f"{self.path_tree}.lock", timeout=60)

check_and_update_all_jobs_status()

Checks the status of all jobs and updates their status in the job dictionary.

This method iterates through all jobs, checks if a ".finished" or a ".failed" file exists in the job's folder, and updates the job's status accordingly. If at least one job is not finished or failed, the overall status is set to "to_finish". If all jobs are finished or failed, the overall status is set to "finished".

Returns:

Type	Description
dict[str, Any]	tuple[dict[str, Any], str]: A tuple containing:
str	A dictionary with all jobs and their updated statuses.
tuple[dict[str, Any], str]	A string representing the final status ("to_finish" or "finished").

Source code in study_da/submit/submit_scan.py

def check_and_update_all_jobs_status(self) -> tuple[dict[str, Any], str]:
    """
    Checks the status of all jobs and updates their status in the job dictionary.

    This method iterates through all jobs, checks if a ".finished" or a ".failed" file exists in
    the job's folder, and updates the job's status accordingly. If at least one job is not
    finished or failed, the overall status is set to "to_finish". If all jobs are finished or
    failed, the overall status is set to "finished".

    Returns:
        tuple[dict[str, Any], str]: A tuple containing:
        - A dictionary with all jobs and their updated statuses.
        - A string representing the final status ("to_finish" or "finished").
    """
    dic_all_jobs = self.get_all_jobs()
    at_least_one_job_to_finish = False
    final_status = "to_finish"
    with self.lock:
        # Get dic tree once to avoid reloading it for every job
        dic_tree = self.dic_tree

        # First pass to update the state of the tree
        for job in dic_all_jobs:
            # Skip jobs that are already finished, failed or unsubmittable
            if nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) in [
                "finished",
                "failed",
                "unsubmittable",
            ]:
                continue

            # Check the state of the others
            relative_job_folder = os.path.dirname(job)
            absolute_job_folder = f"{self.abs_path}/{relative_job_folder}"
            if os.path.exists(f"{absolute_job_folder}/.finished"):
                nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "finished")
            # Check if the job failed otherwise (not to resubmit it again)
            elif os.path.exists(f"{absolute_job_folder}/.failed"):
                nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "failed")
            # else:
            #     at_least_one_job_to_finish = True

        # Second pass to update the state of the tree with unreachable jobs
        dependency_graph = DependencyGraph(dic_tree, dic_all_jobs)
        for job in dic_all_jobs:
            # Get all failed dependencies across the tree
            l_dep_failed = dependency_graph.get_failed_dependency(job)
            if len(l_dep_failed) > 0:
                nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "unsubmittable")
            elif nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) == "to_submit":
                at_least_one_job_to_finish = True

        if not at_least_one_job_to_finish:
            # No more jobs to submit so finished
            dic_tree["status"] = final_status = "finished"
            # Last pass to check if all jobs are properly finished
            for job in dic_all_jobs:
                if nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) != "finished":
                    dic_tree["status"] = final_status = "finished with issues"
                    break

        # Update dic_tree from cluster_submission
        self.dic_tree = dic_tree

    return dic_all_jobs, final_status

configure_jobs(force_configure=False, dic_config_jobs=None)

Configures the jobs by modifying the tree structure and creating the run files for each job.

Parameters:

Name	Type	Description	Default
force_configure	bool	Whether to force reconfiguration. Defaults to False.	False
dic_config_jobs	Optional[dict[str, dict[str, Any]]]	A dictionary containing the configuration of the jobs. Defaults to None.	None

Source code in study_da/submit/submit_scan.py

def configure_jobs(
    self,
    force_configure: bool = False,
    dic_config_jobs: Optional[dict[str, dict[str, Any]]] = None,
) -> None:
    """
    Configures the jobs by modifying the tree structure and creating the run files for each job.

    Args:
        force_configure (bool, optional): Whether to force reconfiguration. Defaults to False.
        dic_config_jobs (Optional[dict[str, dict[str, Any]]], optional): A dictionary containing
            the configuration of the jobs. Defaults to None.
    """
    # Lock since we are modifying the tree
    logging.info("Acquiring lock to configure jobs")
    with self.lock:
        # Get the tree
        dic_tree = self.dic_tree

        # Ensure jobs have not been configured already
        if ("configured" in dic_tree and dic_tree["configured"]) and not force_configure:
            logging.warning("Jobs have already been configured. Skipping.")
            return

        # Configure the jobs (add generation and job keys, set status to "To finish")
        dic_tree = ConfigJobs(dic_tree,starting_depth=-len(Path(self.path_tree).parts) + 2).find_and_configure_jobs(dic_config_jobs)

        # Add the python environment, container image and absolute path of the study to the tree
        dic_tree["python_environment"] = self.path_python_environment
        dic_tree["container_image"] = self.path_container_image
        dic_tree["absolute_path"] = self.abs_path
        dic_tree["status"] = "to_finish"
        dic_tree["configured"] = True

        # Explicitly set the dic_tree property to force rewrite
        self.dic_tree = dic_tree

    logging.info("Jobs have been configured. Lock released.")

generate_run_files(dic_tree, l_jobs, dic_additional_commands_per_gen, dic_dependencies_per_gen, dic_copy_back_per_gen, name_config)

Generates run files for the specified jobs.

Parameters:

Name	Type	Description	Default
dic_tree	dict	The dictionary tree structure.	required
l_jobs	list[str]	List of jobs to submit.	required
dic_additional_commands_per_gen	dict[int, str]	Additional commands per generation. Defaults to {}.	required
dic_dependencies_per_gen	dict[int, list[str]]	Dependencies per generation. Only used when doing a HTC submission.	required
dic_copy_back_per_gen	Optional[dict[int, dict[str, bool]]]	A dictionary containing the files to copy back per generation. Accepted keys are "parquet", "yaml", "txt", "json", "zip" and "all".	required
name_config	str	The name of the configuration file for the study.	required

Returns:

Name	Type	Description
dict	dict	The updated dictionary tree structure.

Source code in study_da/submit/submit_scan.py

def generate_run_files(
    self,
    dic_tree: dict[str, Any],
    l_jobs: list[str],
    dic_additional_commands_per_gen: dict[int, str],
    dic_dependencies_per_gen: dict[int, list[str]],
    dic_copy_back_per_gen: dict[int, dict[str, bool]],
    name_config: str,
) -> dict:
    """
    Generates run files for the specified jobs.

    Args:
        dic_tree (dict): The dictionary tree structure.
        l_jobs (list[str]): List of jobs to submit.
        dic_additional_commands_per_gen (dict[int, str], optional): Additional commands per
            generation. Defaults to {}.
        dic_dependencies_per_gen (dict[int, list[str]], optional): Dependencies per generation.
            Only used when doing a HTC submission.
        dic_copy_back_per_gen (Optional[dict[int, dict[str, bool]]], optional): A dictionary
            containing the files to copy back per generation. Accepted keys are "parquet",
            "yaml", "txt", "json", "zip" and "all".
        name_config (str, optional): The name of the configuration file for the study.

    Returns:
        dict: The updated dictionary tree structure.
    """

    logging.info("Generating run files for the jobs to submit")
    # Generate the run files for the jobs to submit
    dic_all_jobs = self.get_all_jobs()
    for job in l_jobs:
        l_keys = dic_all_jobs[job]["l_keys"]
        job_name = os.path.basename(job)
        relative_job_folder = os.path.dirname(job)
        absolute_job_folder = f"{self.abs_path}/{relative_job_folder}"
        generation_number = dic_all_jobs[job]["gen"]
        submission_type = nested_get(dic_tree, l_keys + ["submission_type"])
        singularity = "docker" in submission_type
        path_python_environment = (
            self.path_python_environment_container
            if singularity
            else self.path_python_environment
        )

        # Ensure that the run file does not already exist
        if "path_run" in nested_get(dic_tree, l_keys):
            path_run_curr = nested_get(dic_tree, l_keys + ["path_run"])
            if path_run_curr is not None and os.path.exists(path_run_curr):
                logging.info(f"Run file already exists for job {job}. Skipping.")
                continue

        # Build l_dependencies and add to the kwargs
        l_dependencies = dic_dependencies_per_gen.get(generation_number, [])

        # Get arguments of current generation
        dic_args = dic_copy_back_per_gen.get(generation_number, {})

        # Mutate the keys
        dic_args = {f"copy_back_{key}": value for key, value in dic_args.items()}

        # Build kwargs for the run file
        kwargs_htc = {
            "l_dependencies": l_dependencies,
            "name_config": name_config,
        } | dic_args

        run_str = generate_run_file(
            absolute_job_folder,
            job_name,
            path_python_environment,
            htc="htc" in submission_type,
            additionnal_command=dic_additional_commands_per_gen.get(generation_number, ""),
            **kwargs_htc,
        )
        # Write the run file
        path_run_job = f"{absolute_job_folder}/run.sh"
        with open(path_run_job, "w") as f:
            f.write(run_str)

        # Change permissions to make the file executable
        os.chmod(path_run_job, 0o755)

        # Record the path to the run file in the tree
        nested_set(dic_tree, l_keys + ["path_run"], path_run_job)

    return dic_tree

get_all_jobs()

Retrieves all jobs from the configuration, without modifying the tree.

Returns:

Name	Type	Description
dict	dict	A dictionary containing all jobs.

Source code in study_da/submit/submit_scan.py

def get_all_jobs(self) -> dict:
    """
    Retrieves all jobs from the configuration, without modifying the tree.

    Returns:
        dict: A dictionary containing all jobs.
    """
    # Get a copy of the tree as it's safer
    with self.lock:
        dic_tree = self.dic_tree
    return ConfigJobs(dic_tree,starting_depth=-len(Path(self.path_tree).parts) + 2).find_all_jobs()

keep_submit_until_done(one_generation_at_a_time=False, wait_time=30, max_try=100, dic_additional_commands_per_gen=None, dic_dependencies_per_gen=None, dic_copy_back_per_gen=None, name_config='config.yaml', force_submit=False)

Keeps submitting jobs until all jobs are finished or failed.

The following arguments are only used for HTC jobs submission: - dic_additional_commands_per_gen - dic_dependencies_per_gen - dic_copy_back_per_gen - name_config

Parameters:

Name	Type	Description	Default
one_generation_at_a_time	bool	Whether to submit one full generation at a time. Defaults to False.	False
wait_time	float	The wait time between submissions in minutes. Defaults to 30.	30
max_try	int	The maximum number of tries before stopping the submission.	100
dic_additional_commands_per_gen	dict[int, str]	Additional commands per generation. Defaults to None.	None
dic_dependencies_per_gen	dict[int, list[str]]	Dependencies per generation. Only used when doing a HTC submission. Defaults to None.	None
dic_copy_back_per_gen	Optional[dict[int, dict[str, bool]]]	A dictionary containing the files to copy back per generation. Accepted keys are "parquet", "yaml", "txt", "json", "zip" and "all". Defaults to None, corresponding to copying back only "light" files, i.e. parquet, yaml and txt.	None
name_config	str	The name of the configuration file for the study. Defaults to "config.yaml".	'config.yaml'
force_submit	bool	If True, jobs are resubmitted even though they failed. Defaults to False.	False

Returns:

Type	Description
None	None

Source code in study_da/submit/submit_scan.py

def keep_submit_until_done(
    self,
    one_generation_at_a_time: bool = False,
    wait_time: float = 30,
    max_try=100,
    dic_additional_commands_per_gen: Optional[dict[int, str]] = None,
    dic_dependencies_per_gen: Optional[dict[int, list[str]]] = None,
    dic_copy_back_per_gen: Optional[dict[int, dict[str, bool]]] = None,
    name_config: str = "config.yaml",
    force_submit: bool = False,
) -> None:
    """
    Keeps submitting jobs until all jobs are finished or failed.

    The following arguments are only used for HTC jobs submission:
    - dic_additional_commands_per_gen
    - dic_dependencies_per_gen
    - dic_copy_back_per_gen
    - name_config

    Args:
        one_generation_at_a_time (bool, optional): Whether to submit one full generation at a
            time. Defaults to False.
        wait_time (float, optional): The wait time between submissions in minutes.
            Defaults to 30.
        max_try (int, optional): The maximum number of tries before stopping the submission.
        dic_additional_commands_per_gen (dict[int, str], optional): Additional commands per
            generation. Defaults to None.
        dic_dependencies_per_gen (dict[int, list[str]], optional): Dependencies per generation.
            Only used when doing a HTC submission. Defaults to None.
        dic_copy_back_per_gen (Optional[dict[int, dict[str, bool]]], optional): A dictionary
            containing the files to copy back per generation. Accepted keys are "parquet",
            "yaml", "txt", "json", "zip" and "all". Defaults to None, corresponding to copying
            back only "light" files, i.e. parquet, yaml and txt.
        name_config (str, optional): The name of the configuration file for the study.
            Defaults to "config.yaml".
        force_submit (bool, optional): If True, jobs are resubmitted even though they failed.
            Defaults to False.


    Returns:
        None
    """
    # Handle mutable default arguments
    if dic_additional_commands_per_gen is None:
        dic_additional_commands_per_gen = {}
    if dic_dependencies_per_gen is None:
        dic_dependencies_per_gen = {}

    if wait_time < 1 / 20:
        logging.warning("Wait time should be at least 10 seconds to prevent locking errors.")
        logging.warning("Setting wait time to 10 seconds.")
        wait_time = 10 / 60

    # I don't need to lock the tree here since the status cheking is read only and
    # the lock is acquired in the submit method for the submission
    while (
        self.submit(
            one_generation_at_a_time,
            dic_additional_commands_per_gen,
            dic_dependencies_per_gen,
            dic_copy_back_per_gen,
            name_config,
            force_submit=force_submit,
        )
        not in ["finished", "finished with issues"]
        and max_try > 0
    ):
        # Wait for a certain amount of time before checking again
        logging.info(f"Waiting {wait_time} minutes before checking again.")
        time.sleep(wait_time * 60)
        max_try -= 1

    if max_try == 0:
        print("Maximum number of tries reached. Stopping submission.")

log_jobs_state(dic_summary_by_gen) staticmethod

Logs the state of jobs for each generation.

Parameters:

Name	Type	Description	Default
dic_summary_by_gen	dict	A dictionary where the keys are generation numbers and the values are dictionaries summarizing job states. Each summary dictionary should contain the following keys: - 'to_submit_later': int, number of jobs left to submit later - 'running_or_queuing': int, number of jobs running or queuing - 'submitted_now': int, number of jobs submitted now - 'finished': int, number of jobs finished - 'failed': int, number of jobs failed - 'dependency_failed': int, number of jobs on hold due to failed dependencies	required

Returns:

Type	Description
None	None

Source code in study_da/submit/submit_scan.py

@staticmethod
def log_jobs_state(dic_summary_by_gen: dict[int, dict[str, int]]) -> None:
    """
    Logs the state of jobs for each generation.

    Args:
        dic_summary_by_gen (dict): A dictionary where the keys are generation numbers
            and the values are dictionaries summarizing job states.
            Each summary dictionary should contain the following keys:
            - 'to_submit_later': int, number of jobs left to submit later
            - 'running_or_queuing': int, number of jobs running or queuing
            - 'submitted_now': int, number of jobs submitted now
            - 'finished': int, number of jobs finished
            - 'failed': int, number of jobs failed
            - 'dependency_failed': int, number of jobs on hold due to failed dependencies

    Returns:
        None
    """
    print("State of the jobs:")
    for gen, dic_summary in dic_summary_by_gen.items():
        print("********************************")
        print(f"Generation {gen}")
        print(f"Jobs left to submit later: {dic_summary['to_submit_later']}")
        print(f"Jobs running or queuing: {dic_summary['running_or_queuing']}")
        print(f"Jobs submitted now: {dic_summary['submitted_now']}")
        print(f"Jobs finished: {dic_summary['finished']}")
        print(f"Jobs failed: {dic_summary['failed']}")
        print(f"Jobs on hold due to failed dependencies: {dic_summary['dependency_failed']}")
        print("********************************")

reset_failed_jobs(dic_tree)

Resets the status of jobs that have failed to "to_submit".

Parameters:

Name	Type	Description	Default
dic_tree	dict[str, Any]	The dictionary tree structure.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: The updated dictionary tree structure.

Source code in study_da/submit/submit_scan.py

def reset_failed_jobs(self, dic_tree: dict[str, Any]) -> dict[str, Any]:
    """
    Resets the status of jobs that have failed to "to_submit".

    Args:
        dic_tree (dict[str, Any]): The dictionary tree structure.

    Returns:
        dict[str, Any]: The updated dictionary tree structure.
    """

    dic_all_jobs = self.get_all_jobs()
    # First pass to update the state of the tree
    for job in dic_all_jobs:
        # Skip jobs that are not failed
        if nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"]) != "failed":
            continue

        # Reset the state of the others
        relative_job_folder = os.path.dirname(job)
        absolute_job_folder = f"{self.abs_path}/{relative_job_folder}"

        # Remove failed tag
        if os.path.exists(f"{absolute_job_folder}/.failed"):
            os.remove(f"{absolute_job_folder}/.failed")
        else:
            logging.warning(f"Failed file not found for job {job}.")

        # Remove run file
        if "path_run" in nested_get(dic_tree, dic_all_jobs[job]["l_keys"]):
            path_run_curr = nested_get(dic_tree, dic_all_jobs[job]["l_keys"] + ["path_run"])
            if path_run_curr is not None and os.path.exists(path_run_curr):
                os.remove(path_run_curr)
            else:
                logging.warning(f"Run file not found for job {job}.")

        # Reset the status of the job
        nested_set(dic_tree, dic_all_jobs[job]["l_keys"] + ["status"], "to_submit")

    return dic_tree

submit(one_generation_at_a_time=False, dic_additional_commands_per_gen=None, dic_dependencies_per_gen=None, dic_copy_back_per_gen=None, name_config='config.yaml', force_submit=False)

Submits the jobs to the cluster. Note that copying back large files (e.g. json colliders) can trigger a throttling mechanism in AFS.

The following arguments are only used for HTC jobs submission: - dic_additional_commands_per_gen - dic_dependencies_per_gen - dic_copy_back_per_gen - name_config

Parameters:

Name	Type	Description	Default
one_generation_at_a_time	bool	Whether to submit one full generation at a time. Defaults to False.	False
dic_additional_commands_per_gen	dict[int, str]	Additional commands per generation. Defaults to None.	None
dic_dependencies_per_gen	dict[int, list[str]]	Dependencies per generation. Only used when doing a HTC submission. Defaults to None.	None
dic_copy_back_per_gen	Optional[dict[int, dict[str, bool]]]	A dictionary containing the files to copy back per generation. Accepted keys are "parquet", "yaml", "txt", "json", "zip" and "all". Defaults to None, corresponding to copying back only "light" files, i.e. parquet, yaml and txt.	None
name_config	str	The name of the configuration file for the study. Defaults to "config.yaml".	'config.yaml'
force_submit	bool	If True, jobs are resubmitted even though they failed. Defaults to False.	False

Returns:

Name	Type	Description
str	str	The final status of the jobs.

Source code in study_da/submit/submit_scan.py

def submit(
    self,
    one_generation_at_a_time: bool = False,
    dic_additional_commands_per_gen: Optional[dict[int, str]] = None,
    dic_dependencies_per_gen: Optional[dict[int, list[str]]] = None,
    dic_copy_back_per_gen: Optional[dict[int, dict[str, bool]]] = None,
    name_config: str = "config.yaml",
    force_submit: bool = False,
) -> str:
    """
    Submits the jobs to the cluster. Note that copying back large files (e.g. json colliders)
    can trigger a throttling mechanism in AFS.

    The following arguments are only used for HTC jobs submission:
    - dic_additional_commands_per_gen
    - dic_dependencies_per_gen
    - dic_copy_back_per_gen
    - name_config

    Args:
        one_generation_at_a_time (bool, optional): Whether to submit one full generation at a
            time. Defaults to False.
        dic_additional_commands_per_gen (dict[int, str], optional): Additional commands per
            generation. Defaults to None.
        dic_dependencies_per_gen (dict[int, list[str]], optional): Dependencies per generation.
            Only used when doing a HTC submission. Defaults to None.
        dic_copy_back_per_gen (Optional[dict[int, dict[str, bool]]], optional): A dictionary
            containing the files to copy back per generation. Accepted keys are "parquet",
            "yaml", "txt", "json", "zip" and "all". Defaults to None, corresponding to copying
            back only "light" files, i.e. parquet, yaml and txt.
        name_config (str, optional): The name of the configuration file for the study.
            Defaults to "config.yaml".
        force_submit (bool, optional): If True, jobs are resubmitted even though they failed.
            Defaults to False.

    Returns:
        str: The final status of the jobs.
    """
    # Handle mutable default arguments
    if dic_additional_commands_per_gen is None:
        dic_additional_commands_per_gen = {}
    if dic_dependencies_per_gen is None:
        dic_dependencies_per_gen = {}
    if dic_copy_back_per_gen is None:
        dic_copy_back_per_gen = {}

    # Handle force submit
    if force_submit:
        logging.warning("Forcing resubmission of all failed jobs.")
        with self.lock:
            # Acquire tree from disk
            dic_tree = self.dic_tree

            # Reset the tree by deleting the failed tags
            dic_tree = self.reset_failed_jobs(dic_tree)
            dic_tree["status"] = "to_finish"
            # Write the tree back to disk
            self.dic_tree = dic_tree

    # Update the status of all jobs before submitting
    dic_all_jobs, final_status = self.check_and_update_all_jobs_status()
    if final_status == "finished":
        print("All jobs are finished.")
        return final_status
    elif final_status == "finished with issues":
        print("All jobs are finished but some did not run properly.")
        return final_status

    logging.info("Acquiring lock to submit jobs")
    with self.lock:
        # Get dic tree once to avoid reloading it for every job
        dic_tree = self.dic_tree

        # Submit the jobs
        self._submit(
            dic_tree,
            dic_all_jobs,
            one_generation_at_a_time,
            dic_additional_commands_per_gen,
            dic_dependencies_per_gen,
            dic_copy_back_per_gen,
            name_config,
        )

        # Update dic_tree from cluster_submission
        self.dic_tree = dic_tree
    logging.info("Jobs have been submitted. Lock released.")
    return final_status

aggregate_output_data(path_tree, l_group_by_parameters, function_to_aggregate=min, generation_of_interest=2, name_output='output_particles.parquet', write_output=True, path_output=None, only_keep_lost_particles=True, dic_parameters_of_interest=None, l_parameters_to_keep=None, name_template_parameters='parameters_lhc.yaml', path_template_parameters=None, force_overwrite=False)

Aggregates output data from simulation files.

Parameters:

Name	Type	Description	Default
path_tree	str	The path to the tree file.	required
l_group_by_parameters	list	List of parameters to group by.	required
function_to_aggregate	callable	Function to aggregate the grouped data.	min
generation_of_interest	int	The generation of interest. Defaults to 2.	2
name_output	str	The name of the output file. Defaults to "output_particles.parquet".	'output_particles.parquet'
write_output	bool	Flag to indicate if the output should be written to a file. Defaults to True.	True
path_output	str	The path to the output file. If not provided, the default output file will be in the study folder as 'da.parquet'. Defaults to None.	None
only_keep_lost_particles	bool	Flag to indicate if only lost particles should be kept. Defaults to True.	True
dic_parameters_of_interest	dict	Dictionary of parameters of interest. Defaults to None.	None
l_parameters_to_keep	list	List of parameters to keep. Defaults to None.	None
name_template_parameters	str	The name of the template parameters file associating each parameter to a list of keys. Defaults to "parameters_lhc.yaml", which is already contained in the study-da package, and includes the main usual parameters.	'parameters_lhc.yaml'
path_template_parameters	str	The path to the template parameters file. Must be provided if a no template already contained in study-da is provided through the argument name_template_parameters. Defaults to None.	None
force_overwrite	bool	Flag to indicate if the output file should be overwritten if it already exists. Defaults to False.	False

Returns:

Type	Description
DataFrame	pd.DataFrame: The final aggregated DataFrame.

Source code in study_da/postprocess/postprocess.py

def aggregate_output_data(
    path_tree: str,
    l_group_by_parameters: List[str],
    function_to_aggregate: Callable = min,
    generation_of_interest: int = 2,
    name_output: str = "output_particles.parquet",
    write_output: bool = True,
    path_output: Optional[str] = None,
    only_keep_lost_particles: bool = True,
    dic_parameters_of_interest: Optional[Dict[str, List[str]]] = None,
    l_parameters_to_keep: Optional[List[str]] = None,
    name_template_parameters: str = "parameters_lhc.yaml",
    path_template_parameters: Optional[str] = None,
    force_overwrite: bool = False,
) -> pd.DataFrame:
    """
    Aggregates output data from simulation files.

    Args:
        path_tree (str): The path to the tree file.
        l_group_by_parameters (list): List of parameters to group by.
        function_to_aggregate (callable, optional): Function to aggregate the grouped data.
        generation_of_interest (int, optional): The generation of interest. Defaults to 2.
        name_output (str, optional): The name of the output file. Defaults to "output_particles.parquet".
        write_output (bool, optional): Flag to indicate if the output should be written to a file.
            Defaults to True.
        path_output (str, optional): The path to the output file. If not provided, the default
            output file will be in the study folder as 'da.parquet'. Defaults to None.
        only_keep_lost_particles (bool, optional): Flag to indicate if only lost particles should be
            kept. Defaults to True.
        dic_parameters_of_interest (dict, optional): Dictionary of parameters of interest. Defaults
            to None.
        l_parameters_to_keep (list, optional): List of parameters to keep. Defaults to None.
        name_template_parameters (str, optional): The name of the template parameters file
            associating each parameter to a list of keys. Defaults to "parameters_lhc.yaml", which
            is already contained in the study-da package, and includes the main usual parameters.
        path_template_parameters (str, optional): The path to the template parameters file. Must
            be provided if a no template already contained in study-da is provided through the
            argument name_template_parameters. Defaults to None.
        force_overwrite (bool, optional): Flag to indicate if the output file should be overwritten
            if it already exists. Defaults to False.

    Returns:
        pd.DataFrame: The final aggregated DataFrame.
    """
    # Check it the output doesn't already exist and ask for confirmation to overwrite
    dic_tree, _ = load_dic_from_path(path_tree)
    absolute_path_study = dic_tree["absolute_path"]
    if path_output is None:
        path_output = os.path.join(absolute_path_study, "da.parquet")
    if os.path.exists(path_output) and not force_overwrite:
        input_user = input(
            f"The output file {path_output} already exists. Do you want to overwrite it? (y/n) "
        )
        if input_user.lower() != "y":
            logging.warning("Output file not overwritten")
            return pd.read_parquet(path_output)

    logging.info("Analysis of output simulation files started")

    dic_all_jobs = ConfigJobs(dic_tree,starting_depth=-len(Path(path_tree).parts) + 2).find_all_jobs()

    l_df_sim = get_particles_data(
        dic_all_jobs, absolute_path_study, generation_of_interest, name_output
    )

    default_path_template_parameters = False
    if dic_parameters_of_interest is None:
        if path_template_parameters is not None:
            logging.info("Loading parameters of interest from the provided configuration file")
        else:
            if name_template_parameters is None:
                raise ValueError(
                    "No template configuration file provided for the parameters of interest"
                )
            logging.info("Loading parameters of interest from the template configuration file")
            path_template_parameters = os.path.join(
                os.path.dirname(inspect.getfile(aggregate_output_data)),
                "configs",
                name_template_parameters,
            )
            default_path_template_parameters = True
        dic_parameters_of_interest, _ = load_dic_from_path(path_template_parameters)

    l_df_output = add_parameters_from_config(
        l_df_sim, dic_parameters_of_interest, default_path_template_parameters
    )

    df_final = merge_and_group_by_parameters_of_interest(
        l_df_output,
        l_group_by_parameters,
        only_keep_lost_particles,
        l_parameters_to_keep,
        function_to_aggregate,
    )

    # Fix the LHC version type
    df_final = fix_LHC_version(df_final)

    if write_output:
        df_final.to_parquet(path_output)
    elif path_output is not None:
        logging.warning("Output path provided but write_output set to False, no output saved")

    logging.info("Final dataframe for current set of simulations: %s", df_final)
    return df_final

create(path_config_scan='config_scan.yaml', force_overwrite=False, dic_parameter_all_gen=None, dic_parameter_all_gen_naming=None, add_prefix_to_folder_names=False)

Create a study based on the configuration file.

Parameters:

Name	Type	Description	Default
path_config_scan	str	Path to the configuration file for the scan. Defaults to "config_scan.yaml".	'config_scan.yaml'
force_overwrite	bool	Flag to force overwrite the study. Defaults to False.	False
dic_parameter_all_gen	Optional[dict[str, dict[str, Any]]]	Dictionary of parameters for the scan, if not provided through the scan config. Defaults to None.	None
dic_parameter_all_gen_naming	Optional[dict[str, dict[str, Any]]]	Dictionary of parameters for the naming of the scan subfolders, if not provided through the scan config. Defaults to None.	None
add_prefix_to_folder_names	bool	Whether to add a prefix to the folder names. Defaults to False.	False

Returns:

Type	Description
tuple[str, str]	tuple[str, str]: The path to the tree file and the name of the main configuration file.

Source code in study_da/study_da.py

def create(
    path_config_scan: str = "config_scan.yaml",
    force_overwrite: bool = False,
    dic_parameter_all_gen: Optional[dict[str, dict[str, Any]]] = None,
    dic_parameter_all_gen_naming: Optional[dict[str, dict[str, Any]]] = None,
    add_prefix_to_folder_names: bool = False,
) -> tuple[str, str]:
    """
    Create a study based on the configuration file.

    Args:
        path_config_scan (str, optional): Path to the configuration file for the scan.
            Defaults to "config_scan.yaml".
        force_overwrite (bool, optional): Flag to force overwrite the study. Defaults to False.
        dic_parameter_all_gen (Optional[dict[str, dict[str, Any]]], optional): Dictionary of
            parameters for the scan, if not provided through the scan config. Defaults to None.
        dic_parameter_all_gen_naming (Optional[dict[str, dict[str, Any]]], optional): Dictionary of
            parameters for the naming of the scan subfolders, if not provided through the scan
            config. Defaults to None.
        add_prefix_to_folder_names (bool, optional): Whether to add a prefix to the folder names.
            Defaults to False.

    Returns:
        tuple[str, str]: The path to the tree file and the name of the main configuration file.
    """
    logging.info(f"Create study from configuration file: {path_config_scan}")
    study = GenerateScan(path_config=path_config_scan)
    study.create_study(
        force_overwrite=force_overwrite,
        dic_parameter_all_gen=dic_parameter_all_gen,
        dic_parameter_all_gen_naming=dic_parameter_all_gen_naming,
        add_prefix_to_folder_names=add_prefix_to_folder_names,
    )

    # Get variables of interest for the submission
    path_tree = study.path_tree
    name_main_configuration = study.config["dependencies"]["main_configuration"]

    return path_tree, name_main_configuration

create_single_job(name_main_configuration, name_executable_generation_1, name_executable_generation_2=None, name_executable_generation_3=None, name_study='single_job_study', force_overwrite=False)

Create a single job study (not a parametric scan) with the specified configuration and executables. Limited to three generations.

Parameters:

Name	Type	Description	Default
name_main_configuration	str	The name of the main configuration file for the study.	required
name_executable_generation_1	str	The name of the executable for the first generation.	required
name_executable_generation_2	Optional[str]	The name of the executable for the second generation. Defaults to None.	None
name_executable_generation_3	Optional[str]	The name of the executable for the third generation. Defaults to None.	None
name_study	str	The name of the study. Defaults to "single_job_study".	'single_job_study'
force_overwrite	bool	Whether to force overwrite existing files. Defaults to False.	False

Returns:

Name	Type	Description
str	str	The path to the tree file.

Source code in study_da/study_da.py

def create_single_job(
    name_main_configuration: str,
    name_executable_generation_1: str,
    name_executable_generation_2: Optional[str] = None,
    name_executable_generation_3: Optional[str] = None,
    name_study: str = "single_job_study",
    force_overwrite: bool = False,
) -> str:
    """
    Create a single job study (not a parametric scan) with the specified configuration and
    executables. Limited to three generations.

    Args:
        name_main_configuration (str): The name of the main configuration file for the study.
        name_executable_generation_1 (str): The name of the executable for the first generation.
        name_executable_generation_2 (Optional[str], optional): The name of the executable for the
            second generation. Defaults to None.
        name_executable_generation_3 (Optional[str], optional): The name of the executable for the
            third generation. Defaults to None.
        name_study (str, optional): The name of the study. Defaults to "single_job_study".
        force_overwrite (bool, optional): Whether to force overwrite existing files.
            Defaults to False.

    Returns:
        str: The path to the tree file.
    """
    # Generate the scan dictionnary
    dic_scan = {
        "name": name_study,
        "dependencies": {"main_configuration": name_main_configuration},
        "structure": {
            "generation_1": {
                "executable": name_executable_generation_1,
            },
        },
    }

    if name_executable_generation_2 is not None:
        dic_scan["structure"]["generation_2"] = {
            "executable": name_executable_generation_2,
        }

    if name_executable_generation_3 is not None:
        dic_scan["structure"]["generation_3"] = {
            "executable": name_executable_generation_3,
        }

    # Create the study
    logging.info(f"Create single job study: {name_study}")
    study = GenerateScan(dic_scan=dic_scan)
    study.create_study(
        force_overwrite=force_overwrite,
    )

    return study.path_tree

get_title_from_configuration(dataframe_data, ions=False, crossing_type=None, display_LHC_version=True, display_energy=True, display_bunch_index=True, display_CC_crossing=True, display_bunch_intensity=True, display_beta=True, display_crossing_IP_1=True, display_crossing_IP_2=True, display_crossing_IP_5=True, display_crossing_IP_8=True, display_bunch_length=True, display_polarity_IP_2_8=True, display_emittance=True, display_chromaticity=True, display_octupole_intensity=True, display_coupling=True, display_filling_scheme=True, display_horizontal_tune=None, display_vertical_tune=None, display_tune=True, display_luminosity_1=True, display_luminosity_2=True, display_luminosity_5=True, display_luminosity_8=True, display_PU_1=True, display_PU_2=True, display_PU_5=True, display_PU_8=True, display_number_of_turns=False)

Generates a title string from the configuration data.

Parameters:

Name	Type	Description	Default
dataframe_data	DataFrame	The dataframe containing configuration data.	required
ions	bool	Whether the beam is composed of ions. Defaults to False.	False
crossing_type	str	The type of crossing: 'vh' or 'hv'. Defaults to None, meaning it will try to be inferred from the optics file name. Back to 'hv' if not found.	None
display_betx_bety	bool	Whether to display the beta functions. Defaults to True.	required
display_LHC_version	bool	Whether to display the LHC version. Defaults to True.	True
display_energy	bool	Whether to display the energy. Defaults to True.	True
display_bunch_index	bool	Whether to display the bunch index. Defaults to True.	True
display_CC_crossing	bool	Whether to display the CC crossing. Defaults to True.	True
display_bunch_intensity	bool	Whether to display the bunch intensity. Defaults to True.	True
display_beta	bool	Whether to display the beta function. Defaults to True.	True
display_crossing_IP_1	bool	Whether to display the crossing at IP1. Defaults to True.	True
display_crossing_IP_2	bool	Whether to display the crossing at IP2. Defaults to True.	True
display_crossing_IP_5	bool	Whether to display the crossing at IP5. Defaults to True.	True
display_crossing_IP_8	bool	Whether to display the crossing at IP8. Defaults to True.	True
display_bunch_length	bool	Whether to display the bunch length. Defaults to True.	True
display_polarity_IP_2_8	bool	Whether to display the polarity at IP2 and IP8. Defaults to True.	True
display_emittance	bool	Whether to display the emittance. Defaults to True.	True
display_chromaticity	bool	Whether to display the chromaticity. Defaults to True.	True
display_octupole_intensity	bool	Whether to display the octupole intensity. Defaults to True.	True
display_coupling	bool	Whether to display the coupling. Defaults to True.	True
display_filling_scheme	bool	Whether to display the filling scheme. Defaults to True.	True
display_horizontal_tune	bool	Whether to display the horizontal tune. Defaults to None. Takes precedence over display_tune.	None
display_vertical_tune	bool	Whether to display the vertical tune. Defaults to None. Takes precedence over display_tune.	None
display_tune	bool	Whether to display the tune. Defaults to True.	True
display_luminosity_1	bool	Whether to display the luminosity at IP1. Defaults to True.	True
display_luminosity_2	bool	Whether to display the luminosity at IP2. Defaults to True.	True
display_luminosity_5	bool	Whether to display the luminosity at IP5. Defaults to True.	True
display_luminosity_8	bool	Whether to display the luminosity at IP8. Defaults to True.	True
display_PU_1	bool	Whether to display the PU at IP1. Defaults to True.	True
display_PU_2	bool	Whether to display the PU at IP2. Defaults to True.	True
display_PU_5	bool	Whether to display the PU at IP5. Defaults to True.	True
display_PU_8	bool	Whether to display the PU at IP8. Defaults to True.	True
display_number_of_turns	bool	Whether to display the number of turns. Defaults to False.	False

Returns:

Name	Type	Description
str	str	The generated title string.

Source code in study_da/plot/build_title.py

def get_title_from_configuration(
    dataframe_data: pd.DataFrame,
    ions: bool = False,
    crossing_type: Optional[str] = None,
    display_LHC_version: bool = True,
    display_energy: bool = True,
    display_bunch_index: bool = True,
    display_CC_crossing: bool = True,
    display_bunch_intensity: bool = True,
    display_beta: bool = True,
    display_crossing_IP_1: bool = True,
    display_crossing_IP_2: bool = True,
    display_crossing_IP_5: bool = True,
    display_crossing_IP_8: bool = True,
    display_bunch_length: bool = True,
    display_polarity_IP_2_8: bool = True,
    display_emittance: bool = True,
    display_chromaticity: bool = True,
    display_octupole_intensity: bool = True,
    display_coupling: bool = True,
    display_filling_scheme: bool = True,
    display_horizontal_tune: Optional[bool] = None,
    display_vertical_tune: Optional[bool] = None,
    display_tune: bool = True,
    display_luminosity_1: bool = True,
    display_luminosity_2: bool = True,
    display_luminosity_5: bool = True,
    display_luminosity_8: bool = True,
    display_PU_1: bool = True,
    display_PU_2: bool = True,
    display_PU_5: bool = True,
    display_PU_8: bool = True,
    display_number_of_turns=False,
) -> str:
    """
    Generates a title string from the configuration data.

    Args:
        dataframe_data (pd.DataFrame): The dataframe containing configuration data.
        ions (bool, optional): Whether the beam is composed of ions. Defaults to False.
        crossing_type (str, optional): The type of crossing: 'vh' or 'hv'. Defaults to None, meaning
            it will try to be inferred from the optics file name. Back to 'hv' if not found.
        display_betx_bety (bool, optional): Whether to display the beta functions. Defaults to True.
        display_LHC_version (bool, optional): Whether to display the LHC version. Defaults to True.
        display_energy (bool, optional): Whether to display the energy. Defaults to True.
        display_bunch_index (bool, optional): Whether to display the bunch index. Defaults to True.
        display_CC_crossing (bool, optional): Whether to display the CC crossing. Defaults to True.
        display_bunch_intensity (bool, optional): Whether to display the bunch intensity. Defaults
            to True.
        display_beta (bool, optional): Whether to display the beta function. Defaults to True.
        display_crossing_IP_1 (bool, optional): Whether to display the crossing at IP1. Defaults to
            True.
        display_crossing_IP_2 (bool, optional): Whether to display the crossing at IP2. Defaults to
            True.
        display_crossing_IP_5 (bool, optional): Whether to display the crossing at IP5. Defaults to
            True.
        display_crossing_IP_8 (bool, optional): Whether to display the crossing at IP8. Defaults to
            True.
        display_bunch_length (bool, optional): Whether to display the bunch length. Defaults to
            True.
        display_polarity_IP_2_8 (bool, optional): Whether to display the polarity at IP2 and IP8.
            Defaults to True.
        display_emittance (bool, optional): Whether to display the emittance. Defaults to True.
        display_chromaticity (bool, optional): Whether to display the chromaticity.
            Defaults to True.
        display_octupole_intensity (bool, optional): Whether to display the octupole intensity.
            Defaults to True.
        display_coupling (bool, optional): Whether to display the coupling. Defaults to True.
        display_filling_scheme (bool, optional): Whether to display the filling scheme. Defaults to
            True.
        display_horizontal_tune (bool, optional): Whether to display the horizontal tune. Defaults to
            None. Takes precedence over display_tune.
        display_vertical_tune (bool, optional): Whether to display the vertical tune. Defaults to
            None. Takes precedence over display_tune.
        display_tune (bool, optional): Whether to display the tune. Defaults to True.
        display_luminosity_1 (bool, optional): Whether to display the luminosity at IP1. Defaults to
            True.
        display_luminosity_2 (bool, optional): Whether to display the luminosity at IP2. Defaults to
            True.
        display_luminosity_5 (bool, optional): Whether to display the luminosity at IP5. Defaults to
            True.
        display_luminosity_8 (bool, optional): Whether to display the luminosity at IP8. Defaults to
            True.
        display_PU_1 (bool, optional): Whether to display the PU at IP1. Defaults to True.
        display_PU_2 (bool, optional): Whether to display the PU at IP2. Defaults to True.
        display_PU_5 (bool, optional): Whether to display the PU at IP5. Defaults to True.
        display_PU_8 (bool, optional): Whether to display the PU at IP8. Defaults to True.
        display_number_of_turns (bool, optional): Whether to display the number of turns. Defaults to
            False.

    Returns:
        str: The generated title string.
    """

    # Warn about tune definition
    if (
        display_horizontal_tune is not None or display_vertical_tune is not None
    ) and not display_tune:
        logging.warning(
            "You have defined display_horizontal_tune or display_vertical_tune, but not "
            "display_tune. The horizontal and/or vertical tunes will still be displayed."
        )

    # Find out what is the crossing type
    if crossing_type is None:
        crossing_type = get_crossing_type(dataframe_data)

    # Collect all the information to display
    LHC_version_str = get_LHC_version_str(dataframe_data, ions)
    energy_str = get_energy_str(dataframe_data, ions)
    bunch_index_str = get_bunch_index_str(dataframe_data)
    CC_crossing_str = get_CC_crossing_str(dataframe_data)
    bunch_intensity_str = get_bunch_intensity_str(dataframe_data)
    beta_str = get_beta_str(dataframe_data)
    xing_IP1_str, xing_IP5_str = get_crossing_IP_1_5_str(dataframe_data, crossing_type)
    xing_IP2_str, xing_IP8_str = get_crossing_IP_2_8_str(dataframe_data)
    bunch_length_str = get_bunch_length_str(dataframe_data)
    polarity_str = get_polarity_IP_2_8_str(dataframe_data)
    emittance_str = get_normalized_emittance_str(dataframe_data)
    chromaticity_str = get_chromaticity_str(dataframe_data)
    octupole_intensity_str = get_octupole_intensity_str(dataframe_data)
    coupling_str = get_linear_coupling_str(dataframe_data)
    filling_scheme_str = get_filling_scheme_str(dataframe_data)
    tune_str = get_tune_str(dataframe_data, display_horizontal_tune, display_vertical_tune)
    n_turns_str = get_number_of_turns_str(dataframe_data)

    # Collect luminosity and PU strings at each IP
    dic_lumi_PU_str = {
        "with_beam_beam": {"lumi": {}, "PU": {}},
        "without_beam_beam": {"lumi": {}, "PU": {}},
    }
    for beam_beam in ["with_beam_beam", "without_beam_beam"]:
        for ip in [1, 2, 5, 8]:
            dic_lumi_PU_str[beam_beam]["lumi"][ip] = get_luminosity_at_ip_str(
                dataframe_data, ip, beam_beam=True
            )
            dic_lumi_PU_str[beam_beam]["PU"][ip] = get_PU_at_IP_str(
                dataframe_data, ip, beam_beam=True
            )

    def test_if_empty_and_add_period(string: str) -> str:
        """
        Test if a string is empty and add a period if not.

        Args:
            string (str): The string to test.

        Returns:
            str: The string with a period if not empty.
        """
        return f"{string}. " if string != "" else ""

    # Make the final title (order is the same as in the past)
    title = ""
    if display_LHC_version:
        title += test_if_empty_and_add_period(LHC_version_str)
    if display_energy:
        title += test_if_empty_and_add_period(energy_str)
    if display_CC_crossing:
        title += test_if_empty_and_add_period(CC_crossing_str)
    if display_bunch_intensity:
        title += test_if_empty_and_add_period(bunch_intensity_str)
    # Jump to the next line
    title += "\n"
    if display_luminosity_1:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["lumi"][1])
    if display_PU_1:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["PU"][1])
    if display_luminosity_5:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["lumi"][5])
    if display_PU_5:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["PU"][5])
    # Jump to the next line
    title += "\n"
    if display_luminosity_2:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["lumi"][2])
    if display_PU_2:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["PU"][2])
    if display_luminosity_8:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["lumi"][8])
    if display_PU_8:
        title += test_if_empty_and_add_period(dic_lumi_PU_str["with_beam_beam"]["PU"][8])
    # Jump to the next line
    title += "\n"
    if display_beta:
        title += test_if_empty_and_add_period(beta_str)
    if display_polarity_IP_2_8:
        title += test_if_empty_and_add_period(polarity_str)
    if display_bunch_length:
        title += test_if_empty_and_add_period(bunch_length_str)
    # Jump to the next line
    title += "\n"
    if display_crossing_IP_1:
        title += test_if_empty_and_add_period(xing_IP1_str)
    if display_crossing_IP_5:
        title += test_if_empty_and_add_period(xing_IP5_str)
    if display_crossing_IP_2:
        title += test_if_empty_and_add_period(xing_IP2_str)
    if display_crossing_IP_8:
        title += test_if_empty_and_add_period(xing_IP8_str)

    # Jump to the next line
    title += "\n"
    if display_emittance:
        title += test_if_empty_and_add_period(emittance_str)
    if display_chromaticity:
        title += test_if_empty_and_add_period(chromaticity_str)
    if display_octupole_intensity:
        title += test_if_empty_and_add_period(octupole_intensity_str)
    if display_coupling:
        title += test_if_empty_and_add_period(coupling_str)
    if display_tune:
        title += test_if_empty_and_add_period(tune_str)
    # Jump to the next line
    title += "\n"
    if display_filling_scheme:
        title += test_if_empty_and_add_period(filling_scheme_str)
    if display_bunch_index:
        title += test_if_empty_and_add_period(bunch_index_str)
    # Jump to the next line
    if display_number_of_turns:
        title += "\n"
        title += test_if_empty_and_add_period(n_turns_str)

    # Filter final title for empty lines
    title = "\n".join([line for line in title.split("\n") if line.strip() != ""])

    return title

plot_3D(dataframe_data, x_variable, y_variable, z_variable, color_variable, xlabel=None, ylabel=None, z_label=None, title='', vmin=4.5, vmax=7.5, surface_count=30, opacity=0.2, figsize=(1000, 1000), colormap='RdBu', colorbar_title_text='Minimum DA (σ)', display_colormap=False, output_path='output.png', output_path_html='output.html', display_plot=True, dark_theme=False)

Plots a 3D volume rendering from the given dataframe.

Parameters:

Name	Type	Description	Default
dataframe_data	DataFrame	The dataframe containing the data to plot.	required
x_variable	str	The variable to plot on the x-axis.	required
y_variable	str	The variable to plot on the y-axis.	required
z_variable	str	The variable to plot on the z-axis.	required
color_variable	str	The variable to use for the color scale.	required
xlabel	Optional[str]	The label for the x-axis. Defaults to None.	None
ylabel	Optional[str]	The label for the y-axis. Defaults to None.	None
z_label	Optional[str]	The label for the z-axis. Defaults to None.	None
title	str	The title of the plot. Defaults to "".	''
vmin	float	The minimum value for the color scale. Defaults to 4.5.	4.5
vmax	float	The maximum value for the color scale. Defaults to 7.5.	7.5
surface_count	int	The number of surfaces for volume rendering. Defaults to 30.	30
opacity	float	The opacity of the volume rendering. Defaults to 0.2.	0.2
figsize	tuple[float, float]	The size of the figure. Defaults to (1000, 1000).	(1000, 1000)
colormap	str	The colormap to use. Defaults to "RdBu".	'RdBu'
colorbar_title_text	str	The label for the colorbar. Defaults to "Minimum DA (σ)".	'Minimum DA (σ)'
display_colormap	bool	Whether to display the colormap. Defaults to False.	False
output_path	str	The path to save the plot image. Defaults to "output.png".	'output.png'
output_path_html	str	The path to save the plot HTML. Defaults to "output.html".	'output.html'
display_plot	bool	Whether to display the plot. Defaults to True.	True
dark_theme	bool	Whether to use a dark theme. Defaults to False.	False

Returns:

Type	Description
Any	go.Figure: The plotly figure object.

Source code in study_da/plot/plot_study.py

def plot_3D(
    dataframe_data: pd.DataFrame,
    x_variable: str,
    y_variable: str,
    z_variable: str,
    color_variable: str,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    z_label: Optional[str] = None,
    title: str = "",
    vmin: float = 4.5,
    vmax: float = 7.5,
    surface_count: int = 30,
    opacity: float = 0.2,
    figsize: tuple[float, float] = (1000, 1000),
    colormap: str = "RdBu",
    colorbar_title_text: str = "Minimum DA (σ)",
    display_colormap: bool = False,
    output_path: str = "output.png",
    output_path_html: str = "output.html",
    display_plot: bool = True,
    dark_theme: bool = False,
) -> Any:
    """
    Plots a 3D volume rendering from the given dataframe.

    Args:
        dataframe_data (pd.DataFrame): The dataframe containing the data to plot.
        x_variable (str): The variable to plot on the x-axis.
        y_variable (str): The variable to plot on the y-axis.
        z_variable (str): The variable to plot on the z-axis.
        color_variable (str): The variable to use for the color scale.
        xlabel (Optional[str], optional): The label for the x-axis. Defaults to None.
        ylabel (Optional[str], optional): The label for the y-axis. Defaults to None.
        z_label (Optional[str], optional): The label for the z-axis. Defaults to None.
        title (str, optional): The title of the plot. Defaults to "".
        vmin (float, optional): The minimum value for the color scale. Defaults to 4.5.
        vmax (float, optional): The maximum value for the color scale. Defaults to 7.5.
        surface_count (int, optional): The number of surfaces for volume rendering. Defaults to 30.
        opacity (float, optional): The opacity of the volume rendering. Defaults to 0.2.
        figsize (tuple[float, float], optional): The size of the figure. Defaults to (1000, 1000).
        colormap (str, optional): The colormap to use. Defaults to "RdBu".
        colorbar_title_text (str, optional): The label for the colorbar. Defaults to "Minimum DA (σ)".
        display_colormap (bool, optional): Whether to display the colormap. Defaults to False.
        output_path (str, optional): The path to save the plot image. Defaults to "output.png".
        output_path_html (str, optional): The path to save the plot HTML. Defaults to "output.html".
        display_plot (bool, optional): Whether to display the plot. Defaults to True.
        dark_theme (bool, optional): Whether to use a dark theme. Defaults to False.

    Returns:
        go.Figure: The plotly figure object.
    """
    # Check if plotly is installed
    try:
        import plotly.graph_objects as go
    except ImportError as e:
        raise ImportError("Please install plotly to use this function") from e

    X = np.array(dataframe_data[x_variable])
    Y = np.array(dataframe_data[y_variable])
    Z = np.array(dataframe_data[z_variable])
    values = np.array(dataframe_data[color_variable])
    fig = go.Figure(
        data=go.Volume(
            x=X.flatten(),
            y=Y.flatten(),
            z=Z.flatten(),
            value=values.flatten(),
            isomin=vmin,
            isomax=vmax,
            opacity=opacity,  # needs to be small to see through all surfaces
            surface_count=surface_count,  # needs to be a large number for good volume rendering
            colorscale=colormap,
            colorbar_title_text=colorbar_title_text,
        )
    )

    fig.update_layout(
        scene_xaxis_title_text=xlabel,
        scene_yaxis_title_text=ylabel,
        scene_zaxis_title_text=z_label,
        title=title,
    )

    # Get a good initial view, dezoomed
    fig.update_layout(scene_camera=dict(eye=dict(x=1.5, y=1.5, z=1.5)))

    # Center the title
    fig.update_layout(title_x=0.5, title_y=0.9, title_xanchor="center", title_yanchor="top")

    # Specify the width and height of the figure
    fig.update_layout(width=figsize[0], height=figsize[1])

    # Remove margins and padding
    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))

    # Display the colormap
    if not display_colormap:
        fig.update_layout(coloraxis_showscale=False)
        fig.update_traces(showscale=False)
    else:
        # Make colorbar smaller
        fig.update_layout(coloraxis_colorbar=dict(thickness=10, len=0.5))

    # Set the theme
    if dark_theme:
        fig.update_layout(template="plotly_dark")

    # Display/save/return the figure
    if output_path is not None:
        fig.write_image(output_path)

    if output_path_html is not None:
        fig.write_html(output_path_html)

    if display_plot:
        fig.show()

    return fig

plot_heatmap(dataframe_data, horizontal_variable, vertical_variable, color_variable, link=None, position_qr='top-right', plot_contours=True, xlabel=None, ylabel=None, tick_interval=2, round_xticks=None, round_yticks=None, symmetric_missing=False, mask_lower_triangle=False, mask_upper_triangle=False, plot_diagonal_lines=True, shift_diagonal_lines=1, xaxis_ticks_on_top=True, title='', vmin=4.5, vmax=7.5, k_masking=-1, green_contour=6.0, min_level_contours=1, max_level_contours=15, delta_levels_contours=0.5, figsize=None, label_cbar='Minimum DA (' + '$\\sigma$' + ')', colormap='coolwarm_r', style='ggplot', output_path='output.png', display_plot=True, latex_fonts=True, vectorize=False, fill_missing_value_with=None, dpi=300)

Plots a heatmap from the given dataframe.

Parameters:

Name	Type	Description	Default
dataframe_data	DataFrame	The dataframe containing the data to plot.	required
horizontal_variable	str	The variable to plot on the horizontal axis.	required
vertical_variable	str	The variable to plot on the vertical axis.	required
color_variable	str	The variable to use for the color scale.	required
link	Optional[str]	A link to encode in a QR code. Defaults to None.	None
plot_contours	bool	Whether to plot contours. Defaults to True.	True
xlabel	Optional[str]	The label for the x-axis. Defaults to None.	None
ylabel	Optional[str]	The label for the y-axis. Defaults to None.	None
tick_interval	int	The interval for the ticks. Defaults to 2.	2
round_xticks	Optional[int]	The number of decimal places to round the x-ticks to. Defaults to None.	None
round_yticks	Optional[int]	The number of decimal places to round the y-ticks to. Defaults to None.	None
symmetric_missing	bool	Whether to make the matrix symmetric by replacing the lower triangle with the upper triangle. Defaults to False.	False
mask_lower_triangle	bool	Whether to mask the lower triangle. Defaults to False.	False
mask_upper_triangle	bool	Whether to mask the upper triangle. Defaults to False.	False
plot_diagonal_lines	bool	Whether to plot diagonal lines. Defaults to True.	True
shift_diagonal_lines	int	The shift for the diagonal lines. Defaults to 1.	1
xaxis_ticks_on_top	bool	Whether to place the x-axis ticks on top. Defaults to True.	True
title	str	The title of the plot. Defaults to "".	''
vmin	float	The minimum value for the color scale. Defaults to 4.5.	4.5
vmax	float	The maximum value for the color scale. Defaults to 7.5.	7.5
k_masking	int	The k parameter for masking. Defaults to -1.	-1
green_contour	Optional[float]	The value for the green contour line. Defaults to 6.0.	6.0
min_level_contours	float	The minimum level for the contours. Defaults to 1.	1
max_level_contours	float	The maximum level for the contours. Defaults to 15.	15
delta_levels_contours	float	The delta between contour levels. Defaults to 0.5.	0.5
figsize	Optional[tuple[float, float]]	The size of the figure. Defaults to None.	None
label_cbar	str	The label for the colorbar. Defaults to "Minimum DA ($\sigma$)".	'Minimum DA (' + '$\\sigma$' + ')'
colormap	str	The colormap to use. Defaults to "coolwarm_r".	'coolwarm_r'
style	str	The style to use for the plot. Defaults to "ggplot".	'ggplot'
output_path	str	The path to save the plot. Defaults to "output.pdf".	'output.png'
display_plot	bool	Whether to display the plot. Defaults to True.	True
latex_fonts	bool	Whether to use LaTeX fonts. Defaults to True.	True
vectorize	bool	Whether to vectorize the plot. Defaults to False.	False
fill_missing_value_with	Optional[str | float]	The value to fill missing values with. Can be a number or 'interpolate'. Defaults to None.	None
dpi	int	The DPI for the plot. Defaults to 300.	300

Returns:

Type	Description
tuple[Figure, Axes]	tuple[plt.Figure, plt.Axes]: The figure and axes of the plot.

Source code in study_da/plot/plot_study.py

def plot_heatmap(
    dataframe_data: pd.DataFrame,
    horizontal_variable: str,
    vertical_variable: str,
    color_variable: str,
    link: Optional[str] = None,
    position_qr: Optional[str] = "top-right",
    plot_contours: bool = True,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    tick_interval: int = 2,
    round_xticks: Optional[int] = None,
    round_yticks: Optional[int] = None,
    symmetric_missing: bool = False,
    mask_lower_triangle: bool = False,
    mask_upper_triangle: bool = False,
    plot_diagonal_lines: bool = True,
    shift_diagonal_lines: int = 1,
    xaxis_ticks_on_top: bool = True,
    title: str = "",
    vmin: float = 4.5,
    vmax: float = 7.5,
    k_masking: int = -1,
    green_contour: Optional[float] = 6.0,
    min_level_contours: float = 1,
    max_level_contours: float = 15,
    delta_levels_contours: float = 0.5,
    figsize: Optional[tuple[float, float]] = None,
    label_cbar: str = "Minimum DA (" + r"$\sigma$" + ")",
    colormap: str = "coolwarm_r",
    style: str = "ggplot",
    output_path: str = "output.png",
    display_plot: bool = True,
    latex_fonts: bool = True,
    vectorize: bool = False,
    fill_missing_value_with: Optional[str | float] = None,
    dpi=300,
) -> tuple[plt.Figure, plt.Axes]:
    """
    Plots a heatmap from the given dataframe.

    Args:
        dataframe_data (pd.DataFrame): The dataframe containing the data to plot.
        horizontal_variable (str): The variable to plot on the horizontal axis.
        vertical_variable (str): The variable to plot on the vertical axis.
        color_variable (str): The variable to use for the color scale.
        link (Optional[str], optional): A link to encode in a QR code. Defaults to None.
        plot_contours (bool, optional): Whether to plot contours. Defaults to True.
        xlabel (Optional[str], optional): The label for the x-axis. Defaults to None.
        ylabel (Optional[str], optional): The label for the y-axis. Defaults to None.
        tick_interval (int, optional): The interval for the ticks. Defaults to 2.
        round_xticks (Optional[int], optional): The number of decimal places to round the x-ticks to.
            Defaults to None.
        round_yticks (Optional[int], optional): The number of decimal places to round the y-ticks to.
            Defaults to None.
        symmetric_missing (bool, optional): Whether to make the matrix symmetric by replacing the
            lower triangle with the upper triangle. Defaults to False.
        mask_lower_triangle (bool, optional): Whether to mask the lower triangle. Defaults to False.
        mask_upper_triangle (bool, optional): Whether to mask the upper triangle. Defaults to False.
        plot_diagonal_lines (bool, optional): Whether to plot diagonal lines. Defaults to True.
        shift_diagonal_lines (int, optional): The shift for the diagonal lines. Defaults to 1.
        xaxis_ticks_on_top (bool, optional): Whether to place the x-axis ticks on top. Defaults to True.
        title (str, optional): The title of the plot. Defaults to "".
        vmin (float, optional): The minimum value for the color scale. Defaults to 4.5.
        vmax (float, optional): The maximum value for the color scale. Defaults to 7.5.
        k_masking (int, optional): The k parameter for masking. Defaults to -1.
        green_contour (Optional[float], optional): The value for the green contour line. Defaults to 6.0.
        min_level_contours (float, optional): The minimum level for the contours. Defaults to 1.
        max_level_contours (float, optional): The maximum level for the contours. Defaults to 15.
        delta_levels_contours (float, optional): The delta between contour levels. Defaults to 0.5.
        figsize (Optional[tuple[float, float]], optional): The size of the figure. Defaults to None.
        label_cbar (str, optional): The label for the colorbar. Defaults to "Minimum DA ($\sigma$)".
        colormap (str, optional): The colormap to use. Defaults to "coolwarm_r".
        style (str, optional): The style to use for the plot. Defaults to "ggplot".
        output_path (str, optional): The path to save the plot. Defaults to "output.pdf".
        display_plot (bool, optional): Whether to display the plot. Defaults to True.
        latex_fonts (bool, optional): Whether to use LaTeX fonts. Defaults to True.
        vectorize (bool, optional): Whether to vectorize the plot. Defaults to False.
        fill_missing_value_with (Optional[str | float], optional): The value to fill missing values
            with. Can be a number or 'interpolate'. Defaults to None.
        dpi (int, optional): The DPI for the plot. Defaults to 300.

    Returns:
        tuple[plt.Figure, plt.Axes]: The figure and axes of the plot.
    """
    # Use the requested style
    _set_style(style, latex_fonts, vectorize)

    # Get the dataframe to plot
    df_to_plot = dataframe_data.pivot(
        index=vertical_variable, columns=horizontal_variable, values=color_variable
    )

    # Get numpy array from dataframe
    data_array = df_to_plot.to_numpy(dtype=float)

    # Replace NaNs with a value if requested
    if fill_missing_value_with is not None:
        if isinstance(fill_missing_value_with, (int, float)):
            data_array[np.isnan(data_array)] = fill_missing_value_with
        elif fill_missing_value_with == "interpolate":
            # Interpolate missing values with griddata
            x = np.arange(data_array.shape[1])
            y = np.arange(data_array.shape[0])
            xx, yy = np.meshgrid(x, y)
            x = xx[~np.isnan(data_array)]
            y = yy[~np.isnan(data_array)]
            z = data_array[~np.isnan(data_array)]
            data_array = griddata((x, y), z, (xx, yy), method="cubic")

    # Mask the lower or upper triangle (checks are done in the function)
    data_array_masked, mask_main_array = _mask(
        mask_lower_triangle, mask_upper_triangle, data_array, k_masking
    )

    # Define colormap and set NaNs to white
    cmap = matplotlib.colormaps.get_cmap(colormap)
    cmap.set_bad("w")

    # Build heatmap, with inverted y axis
    fig, ax = plt.subplots()
    if figsize is not None:
        fig.set_size_inches(figsize)
    im = ax.imshow(data_array_masked, cmap=cmap, vmin=vmin, vmax=vmax)
    ax.invert_yaxis()

    # Add text annotations
    ax = _add_text_annotation(df_to_plot, data_array, ax, vmin, vmax)

    # Smooth data for contours
    mx = _smooth(data_array, symmetric_missing)

    # Plot contours if requested
    if plot_contours:
        ax = _add_contours(
            ax,
            data_array,
            mx,
            green_contour,
            min_level_contours,
            max_level_contours,
            delta_levels_contours,
            mask_main_array,
        )

    if plot_diagonal_lines:
        # Diagonal lines must be plotted after the contour lines, because of bug in matplotlib
        # Shift might need to be adjusted
        ax = _add_diagonal_lines(ax, shift=shift_diagonal_lines)

    # Define title and axis labels
    ax.set_title(
        title,
        fontsize=10,
    )

    # Set axis labels
    ax = _set_labels(
        ax,
        df_to_plot,
        data_array,
        horizontal_variable,
        vertical_variable,
        xlabel,
        ylabel,
        xaxis_ticks_on_top,
        tick_interval,
        round_xticks,
        round_yticks,
    )

    # Create colorbar
    cbar = ax.figure.colorbar(im, ax=ax, fraction=0.026, pad=0.04)
    cbar.ax.set_ylabel(label_cbar, rotation=90, va="bottom", labelpad=15)

    # Remove potential grid
    plt.grid(visible=None)

    # Add QR code with a link to the topright side (a bit experimental, might need adjustments)
    if link is not None:
        fig = add_QR_code(fig, link, position_qr)

    # Save and potentially display the plot
    if output_path is not None:
        if output_path.endswith(".pdf") and not vectorize:
            raise ValueError("Please set vectorize=True to save as PDF")
        elif not output_path.endswith(".pdf") and vectorize:
            raise ValueError("Please set vectorize=False to save as PNG or JPG")
        plt.savefig(output_path, bbox_inches="tight", dpi=dpi)

    if display_plot:
        plt.show()
    return fig, ax