Making LMDB Datasets (original format, deprecated for ASE LMDBs)#

Storing your data in an LMDB ensures very fast random read speeds for the fastest supported throughput. This was the recommended option for the majority of fairchem use cases, but has since been deprecated for ASE LMDB files

This notebook provides an overview of how to create LMDB datasets to be used with the FAIRChem repo. This tutorial is intended for those who wish to use FAIRChem to train on their own datasets. Those interested in just using FAIRChem data need not worry about these steps as they’ve been automated as part of this download script.

from fairchem.core.preprocessing import AtomsToGraphs
from fairchem.core.datasets import LmdbDataset
import ase.io
from ase.build import bulk
from ase.build import fcc100, add_adsorbate, molecule
from ase.constraints import FixAtoms
from ase.calculators.emt import EMT
from ase.optimize import BFGS
import matplotlib.pyplot as plt
import lmdb
import pickle
from tqdm import tqdm
import torch
import os

Generate toy dataset: Relaxation of CO on Cu#

adslab = fcc100("Cu", size=(2, 2, 3))
ads = molecule("CO")
add_adsorbate(adslab, ads, 3, offset=(1, 1))
cons = FixAtoms(indices=[atom.index for atom in adslab if (atom.tag == 3)])
adslab.set_constraint(cons)
adslab.center(vacuum=13.0, axis=2)
adslab.set_pbc(True)
adslab.set_calculator(EMT())
dyn = BFGS(adslab, trajectory="CuCO_adslab.traj", logfile=None)
dyn.run(fmax=0, steps=1000)
/tmp/ipykernel_2839/901556023.py:8: DeprecationWarning: Please use atoms.calc = calc
  adslab.set_calculator(EMT())
False
raw_data = ase.io.read("CuCO_adslab.traj", ":")
len(raw_data)
1001

Initial Structure to Relaxed Energy/Structure (IS2RE/IS2RS) LMDBs#

IS2RE/IS2RS LMDBs utilize the SinglePointLmdb dataset. This dataset expects the data to be contained in a SINGLE LMDB file. In addition to the attributes defined by AtomsToGraph, the following attributes must be added for the IS2RE/IS2RS tasks:

  • pos_relaxed: Relaxed adslab positions

  • sid: Unique system identifier, arbitrary

  • y_init: Initial adslab energy, formerly Data.y

  • y_relaxed: Relaxed adslab energy

  • tags (optional): 0 - subsurface, 1 - surface, 2 - adsorbate

As a demo, we will use the above generated data to create an IS2R* LMDB file.

Initialize AtomsToGraph feature extractor#

a2g = AtomsToGraphs(
    max_neigh=50,
    radius=6,
    r_energy=True,    # False for test data
    r_forces=True,    # False for test data
    r_distances=False,
    r_fixed=True,
)

Initialize LMDB file#

db = lmdb.open(
    "sample_CuCO.lmdb",
    map_size=1099511627776 * 2,
    subdir=False,
    meminit=False,
    map_async=True,
)

Write data to LMDB#

def read_trajectory_extract_features(a2g, traj_path):
    traj = ase.io.read(traj_path, ":")
    tags = traj[0].get_tags()
    images = [traj[0], traj[-1]]
    data_objects = a2g.convert_all(images, disable_tqdm=True)
    data_objects[0].tags = torch.LongTensor(tags)
    data_objects[1].tags = torch.LongTensor(tags)
    return data_objects
system_paths = ["CuCO_adslab.traj"]
idx = 0

for system in system_paths:
    # Extract Data object
    data_objects = read_trajectory_extract_features(a2g, system)
    initial_struc = data_objects[0]
    relaxed_struc = data_objects[1]

    initial_struc.y_init = initial_struc.y # subtract off reference energy, if applicable
    del initial_struc.y
    initial_struc.y_relaxed = relaxed_struc.y # subtract off reference energy, if applicable
    initial_struc.pos_relaxed = relaxed_struc.pos

    # Filter data if necessary
    # FAIRChem filters adsorption energies > |10| eV

    initial_struc.sid = idx  # arbitrary unique identifier

    # no neighbor edge case check
    if initial_struc.edge_index.shape[1] == 0:
        print("no neighbors", traj_path)
        continue

    # Write to LMDB
    txn = db.begin(write=True)
    txn.put(f"{idx}".encode("ascii"), pickle.dumps(initial_struc, protocol=-1))
    txn.commit()
    db.sync()
    idx += 1

db.close()
dataset = LmdbDataset({"src": "sample_CuCO.lmdb"})
len(dataset)
1
dataset[0]
Data(pos=[14, 3], cell=[1, 3, 3], atomic_numbers=[14], natoms=14, tags=[14], edge_index=[2, 635], cell_offsets=[635, 3], edge_distance_vec=[635, 3], energy=3.9893144106683787, forces=[14, 3], fixed=[14], pos_relaxed=[14, 3], sid=0)

Structure to Energy and Forces (S2EF) LMDBs#

S2EF LMDBs utilize the TrajectoryLmdb dataset. This dataset expects a directory of LMDB files. In addition to the attributes defined by AtomsToGraph, the following attributes must be added for the S2EF task:

  • tags (optional): 0 - subsurface, 1 - surface, 2 - adsorbate

  • fid: Frame index along the trajcetory

  • sid- sid: Unique system identifier, arbitrary

Additionally, a “length” key must be added to each LMDB file.

As a demo, we will use the above generated data to create an S2EF LMDB dataset

os.makedirs("s2ef", exist_ok=True)
db = lmdb.open(
    "s2ef/sample_CuCO.lmdb",
    map_size=1099511627776 * 2,
    subdir=False,
    meminit=False,
    map_async=True,
)
tags = raw_data[0].get_tags()
data_objects = a2g.convert_all(raw_data, disable_tqdm=True)


for fid, data in tqdm(enumerate(data_objects), total=len(data_objects)):
    #assign sid
    data.sid = torch.LongTensor([0])

    #assign fid
    data.fid = torch.LongTensor([fid])

    #assign tags, if available
    data.tags = torch.LongTensor(tags)

    # Filter data if necessary
    # FAIRChem filters adsorption energies > |10| eV and forces > |50| eV/A

    # no neighbor edge case check
    if data.edge_index.shape[1] == 0:
        print("no neighbors", traj_path)
        continue

    txn = db.begin(write=True)
    txn.put(f"{fid}".encode("ascii"), pickle.dumps(data, protocol=-1))
    txn.commit()

txn = db.begin(write=True)
txn.put(f"length".encode("ascii"), pickle.dumps(len(data_objects), protocol=-1))
txn.commit()


db.sync()
db.close()
  0%|          | 0/1001 [00:00<?, ?it/s]
  6%|▌         | 56/1001 [00:00<00:01, 552.78it/s]
 13%|█▎        | 126/1001 [00:00<00:01, 634.20it/s]
 19%|█▉        | 192/1001 [00:00<00:01, 642.24it/s]
 26%|██▌       | 259/1001 [00:00<00:01, 650.62it/s]
 33%|███▎      | 328/1001 [00:00<00:01, 663.47it/s]
 40%|███▉      | 396/1001 [00:00<00:00, 667.70it/s]
 46%|████▋     | 463/1001 [00:00<00:00, 667.94it/s]
 53%|█████▎    | 533/1001 [00:00<00:00, 675.60it/s]
 60%|██████    | 602/1001 [00:00<00:00, 676.94it/s]
 67%|██████▋   | 674/1001 [00:01<00:00, 688.15it/s]
 74%|███████▍  | 745/1001 [00:01<00:00, 693.97it/s]
 81%|████████▏ | 815/1001 [00:01<00:00, 693.90it/s]
 88%|████████▊ | 885/1001 [00:01<00:00, 691.04it/s]
 95%|█████████▌| 955/1001 [00:01<00:00, 693.56it/s]
100%|██████████| 1001/1001 [00:01<00:00, 677.09it/s]

dataset = LmdbDataset({"src": "s2ef/"})
len(dataset)
1001
dataset[0]
Data(pos=[14, 3], cell=[1, 3, 3], atomic_numbers=[14], natoms=14, tags=[14], edge_index=[2, 635], cell_offsets=[635, 3], edge_distance_vec=[635, 3], energy=3.9893144106683787, forces=[14, 3], fixed=[14], sid=[1], fid=[1], id='0_0')

Advanced usage#

LmdbDataset supports multiple LMDB files because the need to highly parallelize the dataset construction process. With FAIRChem’s largest split containing 135M+ frames, the need to parallelize the LMDB generation process for these was necessary. If you find yourself needing to deal with very large datasets we recommend parallelizing this process.

Interacting with the LMDBs#

Below we demonstrate how to interact with an LMDB to extract particular information.

dataset = LmdbDataset({"src": "s2ef/"})
data = dataset[0]
data
Data(pos=[14, 3], cell=[1, 3, 3], atomic_numbers=[14], natoms=14, tags=[14], edge_index=[2, 635], cell_offsets=[635, 3], edge_distance_vec=[635, 3], energy=3.9893144106683787, forces=[14, 3], fixed=[14], sid=[1], fid=[1], id='0_0')
energies = torch.tensor([data.energy for data in dataset])
energies
tensor([3.9893, 3.9835, 3.9784,  ..., 3.9684, 3.9684, 3.9684])
plt.hist(energies, bins = 10)
plt.yscale("log")
plt.xlabel("Energies")
plt.show()
../_images/70ff966b9de81dc742342614c724a59eb856b0be98cff6b7329d09daf084fa31.png