process.py

from argparse import ArgumentParser

import numpy as np

from experiment import Experiment
from utils import build_true_dict

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "-e",
        "--experiments",
        "--expts",
        nargs="+",
        help='the top-level directories generated by MolPAL runs. I.e., the directory with the "data" and "chkpts" subdirectories',
    )
    parser.add_argument("--true-csv", help="a CSV file containing the true scoring data")
    parser.add_argument("--smiles-col", type=int, default=0)
    parser.add_argument("--score-col", type=int, default=1)
    parser.add_argument("--no-title-line", action="store_true", default=False)
    parser.add_argument(
        "--maximize",
        action="store_true",
        default=False,
        help="whether the objective for which you are calculating performance should be maximized.",
    )
    parser.add_argument("-N", type=int, help="the iteration for which to calculate reward")
    parser.add_argument(
        "-k", type=int, help="the number of top scores from which to calculate the reward"
    )
    parser.add_argument("-o", "--output", help="the output filepath")

    args = parser.parse_args()
    args.title_line = not args.no_title_line

    d_smi_score = build_true_dict(
        args.true_csv, args.smiles_col, args.score_col, args.title_line, args.maximize
    )
    true_smis_scores = sorted(d_smi_score.items(), key=lambda kv: kv[1], reverse=True)
    true_top_k = true_smis_scores[: args.k]

    rs = []
    num_acquired = []
    incomplete_experiments = []
    for p in args.experiments:
        e = Experiment(p)
        rs.append(e.calculate_reward(args.N, true_top_k, False, False))

    R = np.array(rs)
    N_a = np.array(num_acquired)

    np.save(args.output, R)
    print(f"saved file to: {args.output}!")