def do_UA_simulation(compound, T_range, properties, n_points, biasing_factor, optimum_matching, prior): print('Starting UA') UA_simulation = RJMC_Simulation(compound, T_range, properties, n_points, 1 * 10**6, 0.0, biasing_factor, optimum_matching) UA_simulation.prepare_data() compound_2CLJ = LennardJones_2C(UA_simulation.M_w) UA_simulation.gen_Tmatrix(prior, compound_2CLJ) UA_simulation.set_initial_state(prior, compound_2CLJ, initial_model='UA') UA_simulation.RJMC_Outerloop(prior, compound_2CLJ) UA_simulation.Report() UA_logp_trace = UA_simulation.logp_trace print('UA complete') return UA_logp_trace
def basic(simulation_params): print(simulation_params['priors']) prior = RJMC_Prior(simulation_params['priors']) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() rjmc_simulator = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['steps'], simulation_params['swap_freq'], simulation_params['biasing_factor'], simulation_params['optimum_matching']) rjmc_simulator.prepare_data() print('Simulation Attributes:', rjmc_simulator.get_attributes()) compound_2CLJ = LennardJones_2C(rjmc_simulator.M_w) rjmc_simulator.optimum_bounds = simulation_params['opt_bounds'] rjmc_simulator.gen_Tmatrix(prior, compound_2CLJ) # print(rjmc_simulator.opt_params_AUA) rjmc_simulator.set_initial_state(prior, compound_2CLJ) rjmc_simulator.RJMC_Outerloop(prior, compound_2CLJ) trace, logp_trace, percent_dev_trace, BAR_trace = rjmc_simulator.Report( USE_BAR=simulation_params['USE_BAR']) rjmc_simulator.write_output(simulation_params['priors'], tag=simulation_params['label'], save_traj=simulation_params['save_traj']) path = 'output/' + simulation_params['compound'] + '/' + simulation_params['properties'] + '/' + simulation_params['compound'] + \ '_' + simulation_params['properties'] + '_' + str(simulation_params['steps']) + '_' + simulation_params['label'] + '_' + str(date.today()) +'/runfile.yaml' with open(path, 'w') as outfile: yaml.dump(simulation_params, outfile, default_flow_style=False)
def main(): parser = argparse.ArgumentParser(description='Process input parameters for RJMC 2CLJQ') parser.add_argument('--compound', '-c', type=str, help='Compound to simulate parameters for', required=True) parser.add_argument('--trange', '-t', type=list, help='Temperature range for data to include', required=False) parser.add_argument('--steps', '-s', type=int, help='Number of MCMC steps', required=True) parser.add_argument('--properties', '-p', type=str, help='Properties to include in computation', required=True) parser.add_argument('--priors', '-r', type=dict, help='Values and types of prior distribution', required=False) parser.add_argument('--optimum_matching', '-o', type=list, help='Whether to use optimum matching in transitions', required=False) parser.add_argument('--number_data_points', '-d', type=int, help='Number of data points to include', required=False) parser.add_argument('--swap_freq', '-f', type=float, help='Frequency of model jump proposal', required=False) parser.add_argument('--biasing_factor', '-b', type=str, help='Biasing factors for each model', required=False) parser.add_argument('--label', '-l', type=str, help='label for files', required=False) parser.add_argument('--save_traj', '-j', type=bool, help='Whether to save trajectory files', required=False) T_range = [0.55, 0.95] n_points = 10 swap_freq = 0.1 biasing_factor = [0, 0, 0] optimum_matching = ['True', 'True'] tag = '' save_traj = False prior_values = { 'epsilon': ['exponential', [0,400]], 'sigma': ['exponential', [0,5]], 'L': ['exponential', [0,3]], 'Q': ['exponential', [0,1]]} args = parser.parse_args() print(args.compound) print(args.biasing_factor) compound = args.compound # T_range=args.trange steps = args.steps properties = args.properties if args.trange is not None: T_range = args.trange if args.number_data_points is not None: n_points = args.number_data_points if args.swap_freq is not None: swap_freq = args.swap_freq if args.biasing_factor is not None: biasing_factor_raw = list(args.biasing_factor.split(',')) biasing_factor = [float(i) for i in biasing_factor_raw] if args.optimum_matching is not None: optimum_matching = args.optimum_matching if args.priors is not None: prior_values = args.priors if args.label is not None: tag = args.label if args.save_traj is not None: save_traj = args.save_traj prior = RJMC_Prior(prior_values) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() mcmc_prior_simulation = RJMC_Simulation(compound, T_range, properties, n_points, 5 * 10**5, 0.0, biasing_factor, optimum_matching) mcmc_prior_simulation.prepare_data() compound_2CLJ = LennardJones_2C(mcmc_prior_simulation.M_w) mcmc_prior_simulation.gen_Tmatrix(prior, compound_2CLJ) mcmc_prior_simulation.set_initial_state(prior, compound_2CLJ, initial_model='AUA+Q') mcmc_prior_simulation.RJMC_Outerloop(prior, compound_2CLJ) mcmc_prior_simulation.Report() prior_values['Q'] = mcmc_prior_simulation.refit_prior('exponential') print(prior_values['Q']) #plt.hist(mcmc_prior_simulation.trace_model_1[:,4],bins=50,density=True) #plt.plot(np.linspace(0,1,num=500),gengamma.pdf(np.linspace(0,1,num=500),*prior_values['Q'][1])) #plt.plot(np.linspace(0,1,num=500),expon.pdf(np.linspace(0,1,num=500),0,400)) #plt.show() print('Refitting Prior for Q') prior = RJMC_Prior(prior_values) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() print(prior.Q_prior_values) ''' plt.hist(mcmc_prior_simulation.trace_model_1[:,4],bins=50,density=True) plt.plot(np.linspace(0,1,num=500),expon.pdf(np.linspace(0,1,num=500),*prior.Q_prior_values)) plt.show() ''' aua_path = 'output/C2H2/rhol+Psat/C2H2_rhol+Psat_2000000_aua_only_2019-10-25/trace/trace.npy' auaq_path = 'output/C2H2/rhol+Psat/C2H2_rhol+Psat_2000000_auaq_only_2019-10-25/trace/trace.npy' aua_max,auaq_max = create_map(aua_path,auaq_path) rjmc_simulator = RJMC_Simulation(compound, T_range, properties, n_points, steps, swap_freq, biasing_factor, optimum_matching) rjmc_simulator.prepare_data() print('Simulation Attributes:', rjmc_simulator.get_attributes()) compound_2CLJ = LennardJones_2C(rjmc_simulator.M_w) rjmc_simulator.gen_Tmatrix(prior, compound_2CLJ) rjmc_simulator.set_initial_state(prior, compound_2CLJ) custom_map = list([list(aua_max),list(auaq_max),list(rjmc_simulator.opt_params_UA)]) rjmc_simulator.opt_params_AUA,rjmc_simulator.opt_params_AUA_Q,rjmc_simulator.opt_params_UA = rjmc_simulator.load_custom_map(custom_map) rjmc_simulator.RJMC_Outerloop(prior, compound_2CLJ) trace, logp_trace, percent_dev_trace,BAR_trace = rjmc_simulator.Report(USE_BAR=True) rjmc_simulator.write_output(prior_values, tag=tag, save_traj=save_traj) print('Finished!')
def main(): T_range = [0.55, 0.95] n_points = 10 swap_freq = 0.1 biasing_factor = [0, 0, 0] optimum_matching = ['True', 'True'] tag = '' save_traj = False prior_values = { 'epsilon': ['exponential', [0, 400]], 'sigma': ['exponential', [0, 5]], 'L': ['exponential', [0, 3]], 'Q': ['exponential', [0, 1]] } args = parse_args() print(args.compound) compound = args.compound steps = args.steps properties = args.properties if args.trange is not None: T_range = args.trange if args.number_data_points is not None: n_points = args.number_data_points if args.swap_freq is not None: swap_freq = args.swap_freq if args.biasing_factor is not None: biasing_factor = args.biasing_factor if args.optimum_matching is not None: optimum_matching = args.optimum_matching if args.priors is not None: prior_values = args.priors if args.label is not None: tag = args.label if args.save_traj is not None: save_traj = args.save_traj prior = RJMC_Prior(prior_values) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() ''' mcmc_prior_simulation = RJMC_Simulation(compound, T_range, properties, n_points, 1 * 10**6, 0.0, biasing_factor, optimum_matching) mcmc_prior_simulation.prepare_data() compound_2CLJ = LennardJones_2C(mcmc_prior_simulation.M_w) mcmc_prior_simulation.gen_Tmatrix(prior, compound_2CLJ) mcmc_prior_simulation.set_initial_state(prior, compound_2CLJ, initial_model='AUA+Q') mcmc_prior_simulation.RJMC_Outerloop(prior, compound_2CLJ) mcmc_prior_simulation.Report() prior_values['Q'] = mcmc_prior_simulation.refit_prior('gamma') print('Refitting Prior for Q') prior = RJMC_Prior(prior_values) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() ''' UA_logp_trace = do_UA_simulation(compound, T_range, properties, n_points, biasing_factor, optimum_matching, prior) AUA_Q_logp_trace = do_AUA_Q_simulation(compound, T_range, properties, n_points, biasing_factor, optimum_matching, prior) AUA_logp_trace = do_AUA_simulation(compound, T_range, properties, n_points, biasing_factor, optimum_matching, prior) UA_biasing_factor, AUA_Q_biasing_factor = compute_biasing_factors( UA_logp_trace, AUA_logp_trace, AUA_Q_logp_trace) biasing_factor = np.asarray([0, AUA_Q_biasing_factor, UA_biasing_factor]) print('Biasing factor', biasing_factor) rjmc_simulator = RJMC_Simulation(compound, T_range, properties, n_points, steps, swap_freq, biasing_factor, optimum_matching) rjmc_simulator.prepare_data() print('Simulation Attributes:', rjmc_simulator.get_attributes()) compound_2CLJ = LennardJones_2C(rjmc_simulator.M_w) rjmc_simulator.gen_Tmatrix(prior, compound_2CLJ) print(rjmc_simulator.opt_params_AUA) rjmc_simulator.set_initial_state(prior, compound_2CLJ) rjmc_simulator.RJMC_Outerloop(prior, compound_2CLJ) trace, logp_trace, percent_dev_trace, BAR_trace = rjmc_simulator.Report( USE_BAR=True) rjmc_simulator.write_output(prior_values, tag=tag, save_traj=save_traj) prob = rjmc_simulator.prob unbias_prob = unbias_simulation(biasing_factor, prob) print('Unbias Prob: ', unbias_prob) unbias_BF = [ unbias_prob[0] / unbias_prob[1], unbias_prob[0] / unbias_prob[2] ] print('Unbiased Bayes factors: ', unbias_BF) BF = rjmc_simulator.BF_BAR BAR_LB = [BF[0][1][0], BF[1][1][0]] BAR_UB = [BF[0][1][1], BF[1][1][1]] BAR_prob_LB = undo_bar(BAR_LB) BAR_prob_UB = undo_bar(BAR_UB) BF_prob = undo_bar([BF[0][0], BF[1][0]]) print('Prob from BAR: ', BF_prob) print('LB prob from BAR: ', BAR_prob_LB) print('UB prob from BAR: ', BAR_prob_UB) unbias_prob_BAR = unbias_simulation(biasing_factor, BF_prob) unbias_prob_BAR_LB = unbias_simulation(biasing_factor, BAR_prob_LB) unbias_prob_BAR_UB = unbias_simulation(biasing_factor, BAR_prob_UB) print('Unbias prob BAR: ', unbias_prob_BAR) print('Unbias prob LB BAR: ', unbias_prob_BAR_LB) print('Unbias prob UB BAR: ', unbias_prob_BAR_UB) unbias_BF_BAR = [ unbias_prob_BAR[0] / unbias_prob_BAR[1], unbias_prob_BAR[0] / unbias_prob_BAR[2] ] unbias_BF_BAR_LB = [ unbias_prob_BAR_LB[0] / unbias_prob_BAR_LB[1], unbias_prob_BAR_LB[0] / unbias_prob_BAR_LB[2] ] unbias_BF_BAR_UB = [ unbias_prob_BAR_UB[0] / unbias_prob_BAR_UB[1], unbias_prob_BAR_UB[0] / unbias_prob_BAR_UB[2] ] print('Unbias BF from BAR: ', unbias_BF_BAR) print('Unbias BF LB from BAR: ', unbias_BF_BAR_LB) print('Unbias BF UB from BAR: ', unbias_BF_BAR_UB) print('Finished!')
def refit_prior(simulation_params): prior = RJMC_Prior(simulation_params['priors']) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() print('Approximating AUA+Q posterior distribution') mcmc_prior_simulation = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['single_simulation_length'], 0.0, [0, 0, 0], simulation_params['optimum_matching']) mcmc_prior_simulation.prepare_data() compound_2CLJ = LennardJones_2C(mcmc_prior_simulation.M_w) mcmc_prior_simulation.optimum_bounds = simulation_params['opt_bounds'] mcmc_prior_simulation.gen_Tmatrix(prior, compound_2CLJ) mcmc_prior_simulation.set_initial_state(prior, compound_2CLJ, initial_model='AUA+Q') mcmc_prior_simulation.RJMC_Outerloop(prior, compound_2CLJ) mcmc_prior_simulation.Report() simulation_params['priors']['Q'] = mcmc_prior_simulation.refit_prior( simulation_params['refit_prior_to']) print('Refitting Q prior') print(simulation_params['priors']['Q']) prior = RJMC_Prior(simulation_params['priors']) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() if simulation_params['use_MAP'] is True: aua_max, auaq_max = create_map(simulation_params['MAP_simulations'][0], simulation_params['MAP_simulations'][1]) rjmc_simulator = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['steps'], simulation_params['swap_freq'], simulation_params['biasing_factor'], simulation_params['optimum_matching']) rjmc_simulator.prepare_data() print('Simulation Attributes:', rjmc_simulator.get_attributes()) compound_2CLJ = LennardJones_2C(rjmc_simulator.M_w) rjmc_simulator.optimum_bounds = simulation_params['opt_bounds'] rjmc_simulator.gen_Tmatrix(prior, compound_2CLJ) rjmc_simulator.set_initial_state(prior, compound_2CLJ) if simulation_params['use_MAP'] is True: custom_map = list([ list(aua_max), list(auaq_max), list(rjmc_simulator.opt_params_UA) ]) rjmc_simulator.opt_params_AUA, rjmc_simulator.opt_params_AUA_Q, rjmc_simulator.opt_params_UA = rjmc_simulator.load_custom_map( custom_map) rjmc_simulator.RJMC_Outerloop(prior, compound_2CLJ) trace, logp_trace, percent_dev_trace, BAR_trace = rjmc_simulator.Report( USE_BAR=simulation_params['USE_BAR']) rjmc_simulator.write_output(simulation_params['priors'], tag=simulation_params['label'], save_traj=simulation_params['save_traj']) path = 'output/' + simulation_params['compound'] + '/' + simulation_params['properties'] + '/' + simulation_params['compound'] + \ '_' + simulation_params['properties'] + '_' + str(simulation_params['steps']) + '_' + simulation_params['label'] + '_' + str(date.today()) +'/runfile.yaml' with open(path, 'w') as outfile: yaml.dump(simulation_params, outfile, default_flow_style=False)
def biasing_factor(simulation_params): if simulation_params['refit_prior'] is True: prior = RJMC_Prior(simulation_params['priors']) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() print('Approximating AUA+Q posterior distribution') mcmc_prior_simulation = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['single_simulation_length'], 0.0, [0, 0, 0], simulation_params['optimum_matching']) mcmc_prior_simulation.prepare_data() compound_2CLJ = LennardJones_2C(mcmc_prior_simulation.M_w) mcmc_prior_simulation.gen_Tmatrix(prior, compound_2CLJ) mcmc_prior_simulation.set_initial_state(prior, compound_2CLJ, initial_model='AUA+Q') mcmc_prior_simulation.RJMC_Outerloop(prior, compound_2CLJ) mcmc_prior_simulation.Report() simulation_params['priors']['Q'] = mcmc_prior_simulation.refit_prior( simulation_params['refit_prior_to']) print('Refitting Q prior') print(simulation_params['priors']['Q']) print('Starting AUA') AUA_simulation = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['biasing_factor_simulation_length'], 0.0, [0, 0, 0], simulation_params['optimum_matching']) AUA_simulation.prepare_data() compound_2CLJ = LennardJones_2C(AUA_simulation.M_w) AUA_simulation.gen_Tmatrix(prior, compound_2CLJ) AUA_simulation.set_initial_state(prior, compound_2CLJ, initial_model='AUA') AUA_simulation.RJMC_Outerloop(prior, compound_2CLJ) AUA_simulation.Report() AUA_logp_trace = AUA_simulation.logp_trace print('AUA Complete!') print('Starting AUA+Q') AUAQ_simulation = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['biasing_factor_simulation_length'], 0.0, [0, 0, 0], simulation_params['optimum_matching']) AUAQ_simulation.prepare_data() compound_2CLJ = LennardJones_2C(AUAQ_simulation.M_w) AUAQ_simulation.gen_Tmatrix(prior, compound_2CLJ) AUAQ_simulation.set_initial_state(prior, compound_2CLJ, initial_model='AUA+Q') AUAQ_simulation.RJMC_Outerloop(prior, compound_2CLJ) AUAQ_simulation.Report() AUAQ_logp_trace = AUAQ_simulation.logp_trace print('AUA Complete!') print('Starting UA') UA_simulation = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['biasing_factor_simulation_length'], 0.0, [0, 0, 0], simulation_params['optimum_matching']) UA_simulation.prepare_data() compound_2CLJ = LennardJones_2C(UA_simulation.M_w) UA_simulation.gen_Tmatrix(prior, compound_2CLJ) UA_simulation.set_initial_state(prior, compound_2CLJ, initial_model='UA') UA_simulation.RJMC_Outerloop(prior, compound_2CLJ) UA_simulation.Report() UA_logp_trace = UA_simulation.logp_trace print('UA Complete!') UA_biasing_factor, AUAQ_biasing_factor = compute_biasing_factors( UA_logp_trace, AUA_logp_trace, AUAQ_logp_trace) biasing_factor = np.asarray([0, AUAQ_biasing_factor, UA_biasing_factor]) print('Biasing factor', biasing_factor) rjmc_simulator = RJMC_Simulation( simulation_params['compound'], simulation_params['trange'], simulation_params['properties'], simulation_params['number_data_points'], simulation_params['steps'], simulation_params['swap_freq'], biasing_factor, simulation_params['optimum_matching']) rjmc_simulator.prepare_data() print('Simulation Attributes:', rjmc_simulator.get_attributes()) compound_2CLJ = LennardJones_2C(rjmc_simulator.M_w) rjmc_simulator.optimum_bounds = simulation_params['opt_bounds'] rjmc_simulator.gen_Tmatrix(prior, compound_2CLJ) # print(rjmc_simulator.opt_params_AUA) rjmc_simulator.set_initial_state(prior, compound_2CLJ) rjmc_simulator.RJMC_Outerloop(prior, compound_2CLJ) trace, logp_trace, percent_dev_trace, BAR_trace = rjmc_simulator.Report( USE_BAR=simulation_params['USE_BAR']) rjmc_simulator.write_output(simulation_params['priors'], tag=simulation_params['label'], save_traj=simulation_params['save_traj']) path = 'output/' + simulation_params['compound'] + '/' + simulation_params['properties'] + '/' + simulation_params['compound'] + \ '_' + simulation_params['properties'] + '_' + str(simulation_params['steps']) + '_' + simulation_params['label'] + '_' + str(date.today()) +'/runfile.yaml' with open(path, 'w') as outfile: yaml.dump(simulation_params, outfile, default_flow_style=False)
def main(): parser = argparse.ArgumentParser(description='Process input parameters for RJMC 2CLJQ') parser.add_argument('--compound', '-c', type=str, help='Compound to simulate parameters for', required=True) parser.add_argument('--trange', '-t', type=list, help='Temperature range for data to include', required=False) parser.add_argument('--steps', '-s', type=int, help='Number of MCMC steps', required=True) parser.add_argument('--properties', '-p', type=str, help='Properties to include in computation', required=True) parser.add_argument('--priors', '-r', type=dict, help='Values and types of prior distribution', required=False) parser.add_argument('--optimum_matching', '-o', type=list, help='Whether to use optimum matching in transitions', required=False) parser.add_argument('--number_data_points', '-d', type=int, help='Number of data points to include', required=False) parser.add_argument('--swap_freq', '-f', type=float, help='Frequency of model jump proposal', required=False) parser.add_argument('--biasing_factor', '-b', type=str, help='Biasing factors for each model', required=False) parser.add_argument('--label', '-l', type=str, help='label for files', required=False) parser.add_argument('--save_traj', '-j', type=bool, help='Whether to save trajectory files', required=False) T_range = [0.55, 0.95] n_points = 10 swap_freq = 0.1 biasing_factor = [0, 0, 0] optimum_matching = ['True', 'True'] tag = '' save_traj = False prior_values = { 'epsilon': ['exponential', [400]], 'sigma': ['exponential', [5]], 'L': ['exponential', [3]], 'Q': ['exponential', [1]]} args = parser.parse_args() print(args.compound) print(args.biasing_factor) compound = args.compound #T_range=args.trange steps = args.steps properties = args.properties if args.trange is not None: T_range = args.trange if args.number_data_points is not None: n_points = args.number_data_points if args.swap_freq is not None: swap_freq = args.swap_freq if args.biasing_factor is not None: biasing_factor_raw = list(args.biasing_factor.split(',')) biasing_factor = [float(i) for i in biasing_factor_raw] if args.optimum_matching is not None: optimum_matching = args.optimum_matching if args.priors is not None: prior_values = args.priors if args.label is not None: tag = args.label if args.save_traj is not None: save_traj = args.save_traj prior = RJMC_Prior(prior_values) prior.epsilon_prior() prior.sigma_prior() prior.L_prior() prior.Q_prior() rjmc_simulation = RJMC_Simulation(compound, T_range, properties, n_points, steps, 0.0, biasing_factor, optimum_matching) rjmc_simulation.prepare_data() compound_2CLJ = LennardJones_2C(rjmc_simulation.M_w) rjmc_simulation.gen_Tmatrix(prior, compound_2CLJ) rjmc_simulation.set_initial_state(prior, compound_2CLJ, initial_model='AUA+Q') rjmc_simulation.RJMC_Outerloop(prior, compound_2CLJ) rjmc_simulation.Report() rjmc_simulation.write_output(prior_values, tag='auaq_only', save_traj=save_traj) rjmc_simulator = RJMC_Simulation(compound, T_range, properties, n_points, steps, 0.0, biasing_factor, optimum_matching) rjmc_simulator.prepare_data() print('Simulation Attributes:', rjmc_simulator.get_attributes()) compound_2CLJ = LennardJones_2C(rjmc_simulator.M_w) rjmc_simulator.gen_Tmatrix(prior, compound_2CLJ) print(rjmc_simulator.opt_params_AUA) rjmc_simulator.set_initial_state(prior, compound_2CLJ,initial_model = 'AUA') rjmc_simulator.RJMC_Outerloop(prior, compound_2CLJ) rjmc_simulator.Report() rjmc_simulator.write_output(prior_values, tag='aua_only', save_traj=save_traj) print('Finished!')