autogen/flaml/default/regret.py

import argparse
from os import path
import pandas as pd


def build_regret(all, baseline):
    all = all[all.columns.intersection(baseline.index)]
    return baseline - all


def write_regret(regret, filename):
    regret.to_csv(filename)


def load_result(filename, task_type, metric):
    df = pd.read_csv(filename)
    df = df.loc[
        (df[metric].notnull()) & (df.type == task_type),
        ["task", "fold", "params", metric],
    ]
    df["params"] = df["params"].apply(lambda x: path.splitext(path.basename(eval(x)["_modeljson"]))[0])
    baseline = df.loc[df["task"] == df["params"], ["task", metric]].groupby("task").mean()[metric]
    df = df.pivot_table(index="params", columns="task", values=metric)
    return df, baseline


def main():
    parser = argparse.ArgumentParser(description="Build a regret matrix.")
    parser.add_argument("--result_csv", help="File of experiment results")
    parser.add_argument("--task_type", help="Type of task")
    parser.add_argument("--metric", help="Metric for calculating regret", default="result")
    parser.add_argument("--output", help="Location to write regret CSV to")
    args = parser.parse_args()

    all, baseline = load_result(args.result_csv, args.task_type, args.metric)
    regret = build_regret(all, baseline)
    write_regret(regret, args.output)


if __name__ == "__main__":
    # execute only if run as a script
    main()