make_plot_optimization.py

import os
import pickle
from typing import List, Tuple

import numpy as np
from gpflow.kernels import Matern52
from gpflow.kernels.linears import Linear

from algorithms.gp_on_real_space import GPonRealSpace
from algorithms.uncertain_rf import UncertainRandomForest
from data.load_dataset import load_dataset
from util.mlflow.constants import (
    AT_RANDOM,
    ESM,
    EVE,
    EVE_DENSITY,
    OBSERVED_Y,
    ONE_HOT,
    STD_Y,
    TRANSFORMER,
)
from util.mlflow.convenience_functions import get_mlflow_results_optimization
from visualization.plot_metric_for_dataset import plot_optimization_task
from visualization.plot_metric_for_uncertainties import plot_uncertainty_optimization


# TODO: refactor into util/postprocessing
def compute_metrics_optimization_results(
    results: dict, datasets: list, algos: list, representations: list, seeds: list
) -> Tuple[dict, dict, dict, dict]:
    minObs_dict = {}
    regret_dict = {}
    meanObs_dict = {}
    lastObs_dict = {}
    for dataset in datasets:
        algo_minObs = {}
        algo_regret = {}
        algo_meanObs = {}
        algo_lastObs = {}
        for a in algos:
            reps_minObs = {}
            reps_regret = {}
            reps_meanObs = {}
            reps_lastObs = {}
            if a == "GPsquared_exponential":
                a = "GPsqexp"
            for rep in representations:
                seed_minObs = []
                seed_regret = []
                seed_meanObs = []
                seed_lastObs = []
                for seed in seeds:
                    _results = results[seed][dataset][a][rep][None][OBSERVED_Y]
                    min_observed = [
                        min(_results[:i]) for i in range(1, len(_results) + 1)
                    ]
                    seed_minObs.append(min_observed)
                    mean_observed = [
                        np.mean(_results[:i]) for i in range(1, len(_results) + 1)
                    ]
                    seed_meanObs.append(mean_observed)
                    last_observed = [_results[i] for i in range(0, len(_results) - 1)]
                    seed_lastObs.append(last_observed)
                    _, Y = load_dataset(
                        dataset, representation=ONE_HOT
                    )  # observations irrespective of representation
                    regret = [
                        np.sum(_results[:i]) - np.min(Y)
                        for i in range(1, len(_results) + 1)
                    ]
                    seed_regret.append(regret)
                reps_minObs[rep] = seed_minObs
                reps_regret[rep] = seed_regret
                reps_meanObs[rep] = seed_meanObs
                reps_lastObs[rep] = seed_lastObs
            algo_minObs[a] = reps_minObs
            algo_regret[a] = reps_regret
            algo_meanObs[a] = reps_meanObs
            algo_lastObs[a] = reps_lastObs
        minObs_dict[dataset] = algo_minObs
        regret_dict[dataset] = algo_regret
        meanObs_dict[dataset] = algo_meanObs
        lastObs_dict[dataset] = algo_lastObs
    return minObs_dict, regret_dict, meanObs_dict, lastObs_dict


def plot_optimization_results(
    datasets: List[str],
    algos: List[str],
    representations: List[str],
    seeds: List[int],
    reference_benchmark_rep: List[str],
    plot_calibration: bool = False,
    cached_results: bool = False,
    savefig=True,
) -> None:
    for representation in representations:
        cache_filename = f"/Users/rcml/protein_regression/results/cache/results_optimization_d={'_'.join(datasets)}_a={'_'.join(algos)}_r={representation}.pkl"
        cache_filename_ref = f"/Users/rcml/protein_regression/results/cache/results_optimization_d={'_'.join(datasets)}_a={reference_benchmark_rep}_r={representation}.pkl"
        cache_filename_rand = f"/Users/rcml/protein_regression/results/cache/results_optimization_d={'_'.join(datasets)}_a={AT_RANDOM}_r={representation}.pkl"
        # ALL ALGO RESULTS:
        if cached_results and os.path.exists(cache_filename):
            with open(cache_filename, "rb") as infile:
                results = pickle.load(infile)
        else:
            results = get_mlflow_results_optimization(
                datasets=datasets,
                algos=algos,
                reps=[representation],
                metrics=[OBSERVED_Y, STD_Y],
                seeds=seeds,
            )
            with open(cache_filename, "wb") as outfile:
                pickle.dump(results, outfile)
        # COMPARATIVE REFERENCE RESULTS
        if cached_results and os.path.exists(cache_filename_ref):
            with open(cache_filename_ref, "rb") as infile:
                reference_results = pickle.load(infile)
        else:
            reference_results = get_mlflow_results_optimization(
                datasets=datasets,
                algos=reference_benchmark_rep,
                reps=[None],
                metrics=[OBSERVED_Y],
            )
            with open(cache_filename_ref, "wb") as outfile:
                pickle.dump(reference_results, outfile)
        # RANDOM REFERENCE:
        if cached_results and os.path.exists(cache_filename_rand):
            with open(cache_filename_rand, "rb") as infile:
                random_reference_results = pickle.load(infile)
        else:
            random_reference_results = get_mlflow_results_optimization(
                datasets=datasets,
                algos=[AT_RANDOM],
                reps=[None],
                metrics=[OBSERVED_Y],
                seeds=seeds,
            )
            with open(cache_filename_rand, "wb") as outfile:
                pickle.dump(random_reference_results, outfile)

        minObs_dict, regret_dict, meanObs_dict, lastObs_dict = (
            compute_metrics_optimization_results(
                results=results,
                datasets=datasets,
                algos=algos,
                representations=[representation],
                seeds=seeds,
            )
        )
        ref_minObs_dict, ref_regret_dict, ref_meanObs_dict, ref_lastObs_dict = (
            compute_metrics_optimization_results(
                results=reference_results,
                datasets=datasets,
                algos=reference_benchmark_rep,
                representations=[None],
                seeds=[None],
            )
        )
        (
            random_minObs_dict,
            random_regret_dict,
            random_meanObs_dict,
            random_lastObs_dict,
        ) = compute_metrics_optimization_results(
            results=random_reference_results,
            datasets=datasets,
            algos=[AT_RANDOM],
            representations=[None],
            seeds=seeds,
        )
        # add reference to results # eve-score baseline
        for benchmark in reference_benchmark_rep:
            minObs_dict[datasets[0]][benchmark] = ref_minObs_dict[datasets[0]].get(
                benchmark
            )
            regret_dict[datasets[0]][benchmark] = ref_regret_dict[datasets[0]].get(
                benchmark
            )
            meanObs_dict[datasets[0]][benchmark] = ref_meanObs_dict[datasets[0]].get(
                benchmark
            )
            lastObs_dict[datasets[0]][benchmark] = ref_lastObs_dict[datasets[0]].get(
                benchmark
            )
        # random baseline
        minObs_dict[datasets[0]][AT_RANDOM] = random_minObs_dict[datasets[0]].get(
            AT_RANDOM
        )
        regret_dict[datasets[0]][AT_RANDOM] = random_regret_dict[datasets[0]].get(
            AT_RANDOM
        )
        meanObs_dict[datasets[0]][AT_RANDOM] = random_meanObs_dict[datasets[0]].get(
            AT_RANDOM
        )
        lastObs_dict[datasets[0]][AT_RANDOM] = random_lastObs_dict[datasets[0]].get(
            AT_RANDOM
        )

        plot_optimization_task(
            metric_values=minObs_dict,
            name=f"Best_observed",
            representation=representation,
            dataset=datasets,
            savefig=savefig,
        )
        plot_optimization_task(
            metric_values=regret_dict,
            name=f"Regret",
            representation=representation,
            dataset=datasets,
            legend=True,
            savefig=savefig,
        )
        plot_optimization_task(
            metric_values=meanObs_dict,
            name=f"Mean_observed",
            representation=representation,
            dataset=datasets,
            savefig=savefig,
        )
        plot_optimization_task(
            metric_values=lastObs_dict,
            name=f"Last_observed",
            representation=representation,
            dataset=datasets,
            legend=True,
            savefig=savefig,
        )

        if plot_calibration:
            plot_uncertainty_optimization(
                dataset=datasets[0],
                algos=algos,
                rep=representation,
                seeds=seeds,
                number_quantiles=10,
                stepsize=2,
                min_obs_metrics=minObs_dict,
                regret_metrics=regret_dict,
            )


if __name__ == "__main__":
    # gathers all our results and saves them into a numpy array
    datasets = ["1FQG", "UBQT"]
    representations = [TRANSFORMER, ESM, ONE_HOT, EVE]
    plot_calibration = False
    seeds = [
        11,
        42,
        123,
        54,
        2345,
        987,
        6538,
        78543,
        3465,
        43245,
    ]  # 11, 42, 123, 54, 2345, 987, 6538, 78543, 3465, 43245
    reference_benchmark_rep = [EVE_DENSITY]  # option: VAE_DENSITY

    algos = [
        GPonRealSpace(kernel_factory=lambda: Matern52()).get_name(),
        GPonRealSpace(kernel_factory=lambda: Linear()).get_name(),
        UncertainRandomForest().get_name(),
    ]

    for dataset in datasets:
        plot_optimization_results(
            [dataset],
            algos,
            representations,
            seeds,
            reference_benchmark_rep,
            plot_calibration=plot_calibration,
            cached_results=True,
        )  # TODO: UBQT rep: EVE missing!