run_experiments.py

"""Contains the tools needed to reproduce all major experiments
from Parkinson et al. 2023."""
import argparse
import os
import sys

from core_experiments.constants import protein_constants
from core_experiments.mol_prop_experiments import molfit
from core_experiments.prot_engineering_exp import protein_tune, protein_fit, uncertainty_calibration
from core_experiments.prot_engineering_exp import active_learning
from core_experiments.comp_efficiency_experiments import run_approx_nmll_tests, run_cg_tests
from core_experiments.comp_efficiency_experiments import run_fitcomp_tests, run_lbfgs_tests
from core_experiments.comp_efficiency_experiments import run_optimizer_tests


class ReconfigParser(argparse.ArgumentParser):
    """Reconfigure argparse's parser so an automatic error message
    is generated if no args supplied."""
    def error(self, message):
        self.print_help()
        sys.exit(2)


def gen_arg_parser():
    """Build the command line arg parser."""
    parser = ReconfigParser(description="Use this command line app to "
                "run the experiments described in the paper and reproduce key "
                "results. Output is printed to the console and written to "
                "the log files under final_results.")
    parser.add_argument("test", nargs = 1, help=
            "Test sequence to run. Must be one of: protein_tune,"
                        "protein_fit, molfit, optimizers, "
                        "cg, lbfgs, nmll, fitcomp, uncert_calib, "
                        "active_learn.")
    parser.add_argument("temp_dir", nargs = 1, help=
            "A temporary directory where features generated by the convolution kernels "
            "can be stored.")
    return parser



def prep_exp_logfiles():
    """Sets up the logfiles to which experiment results will
    be written (if they are not already present)."""
    os.chdir("final_results")
    if "conv_tune.txt" not in os.listdir():
        with open("conv_tune.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,Hyperparams,"
                    "Num_Train_RFFs,Tuning_Time,"
                    "NMLL,random_seed\n")
    if "conv_fit.txt" not in os.listdir():
        with open("conv_fit.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,"
                    "Hyperparams,"
                    "Num_Train_RFFs,Num_Fitting_RFFs,"
                    "Fitting_Time,Spearmanr,"
                    "random_seed\n")
    if "moltune.txt" not in os.listdir():
        with open("moltune.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,Hyperparams,"
                    "Num_Train_RFFs,Tuning_Time,"
                    "NMLL\n")
    if "mol_pretune.txt" not in os.listdir():
        with open("mol_pretune.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,Hyperparams,"
                    "Num_Train_RFFs,Tuning_Time,"
                    "NMLL\n")
    if "molfit.txt" not in os.listdir():
        with open("molfit.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,Hyperparams,"
                    "Num_Fit_RFFs,Fitting_Time,"
                    "Test_MAE\n")
    if "optimizelog.txt" not in os.listdir():
        with open("optimizelog.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,"
                    "Hyperparams,"
                    "Num_Train_RFFs,"
                    "Tuning_Time,Num_func_evals,"
                    "optimizer,marginal_negloglik\n")
    if "cglog.txt" not in os.listdir():
        with open("cglog.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Algorithm,Dataset,Kernel_Type,"
                    "Hyperparams,"
                    "Num_Fitting_RFFs,"
                    "Fitting_Time,Niter,"
                    "achieved_ratio,max_rank,method\n")
    if "fitcomp_log.txt" not in os.listdir():
        with open("fitcomp_log.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,"
                    "Hyperparams,"
                    "Num_Fitting_RFFs,"
                    "Method,P_ratio,max_rank,"
                    "Losses\n")
    if "approx_nmll_log.txt" not in os.listdir():
        with open("approx_nmll_log.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,"
                    "Hyperparams,"
                    "Num_Training_RFFs,"
                    "exact_nmll,approx_nmll,"
                    "condition_number\n")
    if "uncert_quant.txt" not in os.listdir():
        with open("uncert_quant.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Dataset,Kernel_Type,"
                    "Hyperparams,"
                    "Num_Fitting_RFFs,"
                    "MAE,spearman_r,"
                    "miscalibration_area\n")
    if "active_learn_log.txt" not in os.listdir():
        with open("active_learn_log.txt", "w+", encoding="utf8") as logfile:
            logfile.write("Random_seed,Iteration,Mean_fitness,Max_fitness\n")
    os.chdir("..")




def main():
    """Entry point for all key experiments."""
    start_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(start_dir)
    parser = gen_arg_parser()
    if len(sys.argv) < 1:
        parser.print_help(sys.stderr)
        sys.exit(1)
    args = parser.parse_args()
    prep_exp_logfiles()

    task = args.test[0]
    if task == "protein_tune":
        protein_tune(start_dir, protein_constants.conv_tuning_params,
                "conv_tune.txt")
    elif task == "optimizers":
        run_optimizer_tests(start_dir)
    elif task == "cg":
        run_cg_tests(start_dir)
    elif task == "lbfgs":
        run_lbfgs_tests(start_dir)
    elif task == "nmll":
        run_approx_nmll_tests(start_dir)
    elif task == "fitcomp":
        run_fitcomp_tests(start_dir)
    elif task == "protein_fit":
        protein_fit(start_dir, "conv_fit.txt", "conv_tune.txt")
    elif task == "molfit":
        molfit(start_dir, args.temp_dir[0])
    elif task == "uncert_calib":
        uncertainty_calibration(start_dir)
    elif task == "active_learn":
        active_learning(start_dir)


if __name__ == "__main__":
    main()