PySR / benchmarks /print_best_model.py
pre-commit-ci[bot]
[pre-commit.ci] pre-commit autoupdate (#537)
75c23d4 unverified
"""Print the best model parameters and loss"""
import pickle as pkl
from pprint import PrettyPrinter
import hyperopt
import numpy as np
from hyperopt import Trials, fmin, hp, tpe
from space import space
# Change the following code to your file
################################################################################
# TODO: Declare a folder to hold all trials objects
TRIALS_FOLDER = "trials2"
################################################################################
def merge_trials(trials1, trials2_slice):
"""Merge two hyperopt trials objects
:trials1: The primary trials object
:trials2_slice: A slice of the trials object to be merged,
obtained with, e.g., trials2.trials[:10]
:returns: The merged trials object
"""
max_tid = 0
if len(trials1.trials) > 0:
max_tid = max([trial["tid"] for trial in trials1.trials])
for trial in trials2_slice:
tid = trial["tid"] + max_tid + 1
hyperopt_trial = Trials().new_trial_docs(
tids=[None], specs=[None], results=[None], miscs=[None]
)
hyperopt_trial[0] = trial
hyperopt_trial[0]["tid"] = tid
hyperopt_trial[0]["misc"]["tid"] = tid
for key in hyperopt_trial[0]["misc"]["idxs"].keys():
hyperopt_trial[0]["misc"]["idxs"][key] = [tid]
trials1.insert_trial_docs(hyperopt_trial)
trials1.refresh()
return trials1
np.random.seed()
# Load up all runs:
import glob
path = TRIALS_FOLDER + "/*.pkl"
files = 0
for fname in glob.glob(path):
trials_obj = pkl.load(open(fname, "rb"))
n_trials = trials_obj["n"]
trials_obj = trials_obj["trials"]
if files == 0:
trials = trials_obj
else:
trials = merge_trials(trials, trials_obj.trials[-n_trials:])
files += 1
print(files, "trials merged")
best_loss = np.inf
best_trial = None
try:
trials
except NameError:
raise NameError("No trials loaded. Be sure to set the right folder")
# for trial in trials:
# if trial['result']['status'] == 'ok':
# loss = trial['result']['loss']
# if loss < best_loss:
# best_loss = loss
# best_trial = trial
# print(best_loss, best_trial['misc']['vals'])
# trials = sorted(trials, key=lambda x: (x['result']['loss'] if trials['result']['status'] == 'ok' else float('inf')))
clean_trials = []
for trial in trials:
clean_trials.append((trial["result"]["loss"], trial["misc"]["vals"]))
clean_trials = sorted(clean_trials, key=lambda x: x[0])
pp = PrettyPrinter(indent=4)
for trial in clean_trials:
loss, params = trial
for k, value in params.items():
value = value[0]
if isinstance(value, int):
possible_args = space[k].pos_args[1:]
try:
value = possible_args[value].obj
except AttributeError:
value = [arg.obj for arg in possible_args[value].pos_args]
params[k] = value
pp.pprint({"loss": loss, "params": params})