|
|
import os |
|
|
import pandas as pd |
|
|
import random |
|
|
import time |
|
|
|
|
|
from code_generation import generate_and_debug, prepare_working_folder |
|
|
|
|
|
def select_seed_implementations( |
|
|
total_num_sample_solvers, |
|
|
num_sample_for_refine=None, |
|
|
): |
|
|
if ( |
|
|
num_sample_for_refine is None or |
|
|
num_sample_for_refine > total_num_sample_solvers or |
|
|
num_sample_for_refine == -1 |
|
|
): |
|
|
num_sample_for_refine = total_num_sample_solvers |
|
|
|
|
|
|
|
|
selected_indices = random.sample(range(total_num_sample_solvers), num_sample_for_refine) |
|
|
|
|
|
return selected_indices |
|
|
|
|
|
|
|
|
|
|
|
def refine(cfg): |
|
|
num_repeated_samples = cfg.method.num_repeated_samples |
|
|
num_trials = cfg.method.num_debugging_trials_per_sample |
|
|
pde_name = cfg.pde.name |
|
|
working_folder = cfg.working_folder |
|
|
model_name = cfg.model.name |
|
|
num_sample_for_refine = cfg.method.num_sample_for_refine |
|
|
start_round = cfg.method.start_round |
|
|
use_sample_solver_init = cfg.method.use_sample_solver_init |
|
|
assert use_sample_solver_init, 'Sample solvers must be enabled for refinement' |
|
|
|
|
|
sample_solver_folder = os.path.join( |
|
|
'solvers', pde_name, cfg.pde.pde_setting_name, 'seeds' |
|
|
) |
|
|
sample_solver_info = pd.read_csv( |
|
|
os.path.join(sample_solver_folder, 'seed_results.csv') |
|
|
) |
|
|
total_num_sample_solvers = len(sample_solver_info) |
|
|
|
|
|
if start_round == 0: |
|
|
prepare_working_folder( |
|
|
cfg, |
|
|
working_folder=working_folder, |
|
|
pde_name=pde_name, |
|
|
use_sample_solver_init=use_sample_solver_init |
|
|
) |
|
|
|
|
|
for round_idx in range(start_round, num_repeated_samples): |
|
|
try: |
|
|
seed_implementations = select_seed_implementations( |
|
|
total_num_sample_solvers=total_num_sample_solvers, |
|
|
num_sample_for_refine=num_sample_for_refine |
|
|
) |
|
|
generate_and_debug( |
|
|
cfg, |
|
|
round_idx=round_idx, |
|
|
num_trials=num_trials, |
|
|
pde_name=pde_name, |
|
|
working_folder=working_folder, |
|
|
seed_implementations=seed_implementations, |
|
|
model_name=model_name |
|
|
) |
|
|
except Exception as e: |
|
|
print(f'Error in sample {round_idx}: {e}. Move on to the next sample.') |
|
|
|
|
|
time.sleep(2) |
|
|
|
