1import ase.db 2from ase.io.formats import string2index 3 4 5def read_db(filename, index, **kwargs): 6 db = ase.db.connect(filename, serial=True, **kwargs) 7 8 if isinstance(index, str): 9 try: 10 index = string2index(index) 11 except ValueError: 12 pass 13 14 if isinstance(index, int): 15 index = slice(index, index + 1 or None) 16 17 if isinstance(index, str): 18 # index is a database query string: 19 for row in db.select(index): 20 yield row.toatoms() 21 else: 22 start, stop, step = index.indices(db.count()) 23 if start == stop: 24 return 25 assert step == 1 26 for row in db.select(offset=start, limit=stop - start): 27 yield row.toatoms() 28 29 30def write_db(filename, images, append=False, **kwargs): 31 con = ase.db.connect(filename, serial=True, append=append, **kwargs) 32 for atoms in images: 33 con.write(atoms) 34 35 36read_json = read_db 37write_json = write_db 38read_postgresql = read_db 39write_postgresql = write_db 40read_mysql = read_db 41write_mysql = write_db 42