File size: 6,196 Bytes
bf314e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | from __future__ import annotations
import argparse
import glob
import logging
import multiprocessing as mp
import os
import random
from pathlib import Path
import numpy as np
from ase.db import connect
from ase.io import read
from tqdm import tqdm
from onescience.datapipes.materials.custom_stack.storage import AseDBDataset
logging.basicConfig(level=logging.INFO)
def compute_normalizer_and_linear_reference(train_path, num_workers):
"""
Given a path to an ASE database file, compute the normalizer value and linear
reference coefficients. These are used to normalize energies and forces during
training. For large datasets, compute this for only a subset of the data.
"""
global dataset
dataset = AseDBDataset({"src": str(train_path)})
sample_indices = random.sample(range(len(dataset)), min(100000, len(dataset)))
with mp.Pool(num_workers) as pool:
outputs = list(
tqdm(
pool.imap(extract_energy_and_forces, sample_indices),
total=len(sample_indices),
desc="Computing normalizer values.",
)
)
atomic_numbers = [x[0] for x in outputs]
energies = [x[1] for x in outputs]
forces = np.array([force for x in outputs for force in x[2]])
force_rms = np.sqrt(np.mean(np.square(forces)))
assert (
np.isfinite(force_rms).all()
), "We found non-finite values in the forces, please check your input data!"
coeff = compute_lin_ref(atomic_numbers, energies)
return force_rms, coeff
def extract_energy_and_forces(idx):
"""
Extract energy and forces from an ASE atoms object at a given index in the dataset.
"""
atoms = dataset.get_atoms(idx)
atomic_numbers = atoms.get_atomic_numbers()
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
n_atoms = len(atoms)
fixed_idx = np.zeros(n_atoms)
if hasattr(atoms, "constraints"):
from ase.constraints import FixAtoms
for constraint in atoms.constraints:
if isinstance(constraint, FixAtoms):
fixed_idx[constraint.index] = 1
mask = fixed_idx == 0
forces = forces[mask]
return atomic_numbers, energy, forces
def compute_lin_ref(atomic_numbers, energies):
"""
Compute linear reference coefficients given atomic numbers and energies.
"""
features = [np.bincount(x, minlength=100).astype(int) for x in atomic_numbers]
X = np.vstack(features)
y = energies
coeff = np.linalg.lstsq(X, y, rcond=None)[0]
assert np.isfinite(
coeff
).all(), "Found some non-finite values while computing element references, please check/sanitize your inputs before proceeding"
return coeff.tolist()
def write_ase_db(mp_arg):
"""
Write ASE atoms objects to an ASE database file. This function is designed to be
run in parallel using multiprocessing.
"""
db_file, file_list, worker_id = mp_arg
successful = []
failed = []
natoms = []
with connect(str(db_file)) as db:
for file in tqdm(file_list, position=worker_id):
atoms_list = read(file, ":")
for i, atoms in enumerate(atoms_list):
try:
assert (
atoms.calc is not None
), "No calculator attached to atoms object."
assert "energy" in atoms.calc.results, "Missing energy result"
assert "forces" in atoms.calc.results, "Missing forces result"
db.write(atoms, data=atoms.info)
natoms.append(len(atoms))
successful.append(f"{file},{i}")
except AssertionError as e:
failed.append(f"{file},{i}: {e!s}")
return db_file, natoms, successful, failed
def launch_processing(data_dir, output_dir, num_workers):
"""
Driver script to launch processing of ASE atoms files into an ASE database.
"""
os.makedirs(output_dir, exist_ok=True)
input_files = [
f
for f in glob.glob(os.path.join(data_dir, "**/*"), recursive=True)
if os.path.isfile(f)
]
chunked_files = np.array_split(input_files, num_workers)
db_files = [output_dir / f"data.{i:04d}.aselmdb" for i in range(num_workers)]
mp_args = [(db_files[i], chunked_files[i], i) for i in range(num_workers)]
with mp.Pool(num_workers) as pool:
outputs = pool.map(write_ase_db, mp_args)
# Log results
natoms = []
for output in outputs:
db_file, _natoms, successful, failed = output
natoms.extend(_natoms)
log_file = db_file.with_suffix(".log")
failed_file = db_file.with_suffix(".failed")
with open(log_file, "w") as log:
log.write("\n".join(successful))
with open(failed_file, "w") as failed_log:
failed_log.write("\n".join(failed))
np.savez_compressed(output_dir / "metadata.npz", natoms=natoms)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--train-dir",
type=str,
required=True,
help="Directory of ASE atoms objects to convert for training.",
)
parser.add_argument(
"--val-dir",
type=str,
required=True,
help="Directory of ASE atoms objects to convert for validation.",
)
parser.add_argument(
"--output-dir",
type=Path,
required=True,
help="Output directory to save required finetuning artifacts.",
)
parser.add_argument(
"--num-workers",
type=int,
default=8,
help="Number of parallel workers for processing files.",
)
args = parser.parse_args()
# Launch processing for training data
train_path = args.output_dir / "train"
launch_processing(args.train_dir, train_path, args.num_workers)
force_rms, linref_coeff = compute_normalizer_and_linear_reference(
train_path, args.num_workers
)
val_path = args.output_dir / "val"
launch_processing(args.val_dir, val_path, args.num_workers)
logging.info(f"Generated dataset at {args.output_dir}") |