run_benchmark.py

from argparse import ArgumentParser

from benchmark import (
    check_benchmark,
    check_surrogate,
    compare_models,
    get_surrogate,
    run_benchmark,
)
from utils import nice_print, read_yaml_config


def main(args):
    """
    Main function to run the benchmark. It reads the config file, checks the benchmark
    configuration, runs the benchmark for each surrogate model, and compares the models
    if specified in the config file.

    Args:
        args (Namespace): The command line arguments.
    """

    config = read_yaml_config(args.config)
    check_benchmark(config)
    surrogates = config["surrogates"]
    # Create dictionary to store metrics for all surrogate models
    all_metrics = {surrogate: {} for surrogate in surrogates}

    # Run benchmark for each surrogate model
    for surrogate_name in surrogates:
        surrogate_class = get_surrogate(surrogate_name)
        if surrogate_class is not None:
            nice_print(f"Running benchmark for {surrogate_name}")
            check_surrogate(surrogate_name, config)
            metrics = run_benchmark(surrogate_name, surrogate_class, config)
            all_metrics[surrogate_name] = metrics
        else:
            print(f"Surrogate {surrogate_name} not recognized. Skipping.")

    # Compare models
    if config["compare"]:
        if len(surrogates) < 2:
            nice_print("At least two surrogate models are required to compare.")
        else:
            nice_print("Comparing models")
            compare_models(all_metrics, config)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--config", default="config.yaml", type=str, help="Path to the config file."
    )
    args = parser.parse_args()
    main(args)