4
1
mirror of https://github.com/pfloos/quack synced 2024-09-27 12:00:58 +02:00
quack/tests/create_database.py
2024-08-30 20:15:39 +02:00

53 lines
1.3 KiB
Python

import argparse
from molecule import save_molecules_to_json, load_molecules_from_json
from molecule import create_database, add_molecule_to_db, remove_database
from feather_bench import FeatherBench
parser = argparse.ArgumentParser(description="Benchmark Data Sets")
parser.add_argument(
'-s', '--set_type',
choices=['light', 'medium', 'heavy'],
default='light',
help="Specify the type of data set: light (default), medium, or heavy."
)
args = parser.parse_args()
if args.set_type == 'light':
bench = 'FeatherBench'
bench_title = "\n\nSelected Light Benchmark: {}\n\n".format(bench)
elif args.set_type == 'medium':
bench = 'BalanceBench'
bench_title = "\n\nSelected Medium Benchmark: {}\n\n".format(bench)
elif args.set_type == 'heavy':
bench = 'TitanBench'
bench_title = "\n\nSelected Heavy Benchmark: {}\n\n".format(bench)
else:
bench_title = "\n\nSelected Light Benchmark: {}\n\n".format(bench)
db_name = '{}.db'.format(bench)
# Save molecules to JSON
#save_molecules_to_json(FeatherBench, 'FeatherBench.json')
# Load molecules from JSON
#loaded_molecules = load_molecules_from_json('FeatherBench.json')
#print(loaded_molecules)
#remove_database(db_name)
create_database(db_name)
for molecule in FeatherBench:
add_molecule_to_db(db_name, molecule)