make_aselmdb.py

from ase.io import read
from fairchem.core.datasets.lmdb_database import LMDBDatabase
from ase.io import read
from tqdm import tqdm
import os
import sys

mode = 'train'
#ld = sorted(os.listdir(f'extxyz/{mode}/'))
f = open(f'extxyz/{mode}_list')
path = f.readlines()
f.close()

proc = {'train': 24, 'val': 10}[mode]

i_start = proc*int(sys.argv[1])
i_end = min(proc*(int(sys.argv[1])+1), len(path))

for i in range(i_start, i_end):
    atoms_list = read(f'extxyz/{path[i].strip()}', ':')
    spl = path[i].split('/')
    directory = '/'.join(spl[:-1])
    fname = spl[-1].replace('extxyz', 'aselmdb').strip()
    fname = fname.replace('atoms', 'omol25')
    print(directory, fname)
    os.makedirs(f'aselmdb/{directory}', exist_ok=True)
    
    with LMDBDatabase(f'aselmdb/{directory}/{fname}') as db:
        for atoms in tqdm(atoms_list):
            db.write(atoms, data=atoms.info)