File size: 4,120 Bytes
fecdc11 | 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 | """Build reproducible ASE DB splits from FairChem's oxide tutorial data."""
from __future__ import annotations
import argparse
import json
import random
from collections import Counter
from pathlib import Path
import numpy as np
from ase import Atoms, units
from ase.calculators.singlepoint import SinglePointCalculator
from ase.db import connect
def load_records(path: Path):
with path.open() as handle:
source = json.load(handle)
records = []
for oxide, polymorphs in source.items():
for polymorph, calculations in polymorphs.items():
for calculation in calculations.get("PBE", {}).get("EOS", {}).get(
"calculations", []
):
records.append((oxide, polymorph, calculation))
return records
def split_groups(records, seed: int):
groups = sorted({(oxide, polymorph) for oxide, polymorph, _ in records})
random.Random(seed).shuffle(groups)
n_train = int(0.8 * len(groups))
n_val = int(0.1 * len(groups))
split_for = {group: "train" for group in groups[:n_train]}
split_for.update({group: "val" for group in groups[n_train : n_train + n_val]})
split_for.update({group: "test" for group in groups[n_train + n_val :]})
return split_for
def to_atoms(oxide: str, polymorph: str, calculation: dict) -> Atoms:
structure = calculation["atoms"]
results = calculation["data"]
atoms = Atoms(
symbols=structure["symbols"],
positions=structure["positions"],
cell=structure["cell"],
pbc=structure["pbc"],
)
atoms.set_tags(np.ones(len(atoms), dtype=int))
# ASE stress uses eV/Angstrom^3. The source JSON documents stress in GPa.
stress = np.asarray(results["stress"], dtype=float) * units.GPa
atoms.calc = SinglePointCalculator(
atoms,
energy=float(results["total_energy"]),
forces=np.asarray(results["forces"], dtype=float),
stress=stress,
)
atoms.info.update({"oxide": oxide, "polymorph": polymorph, "xc": "PBE"})
return atoms
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--input", required=True, type=Path)
parser.add_argument("--output", required=True, type=Path)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument(
"--overwrite", action="store_true", help="Replace existing split databases"
)
args = parser.parse_args()
records = load_records(args.input)
if not records:
raise ValueError(f"No PBE EOS structures found in {args.input}")
args.output.mkdir(parents=True, exist_ok=True)
paths = {split: args.output / f"{split}.db" for split in ("train", "val", "test")}
existing = [str(path) for path in paths.values() if path.exists()]
if existing and not args.overwrite:
raise FileExistsError("Output exists; pass --overwrite to replace: " + ", ".join(existing))
for path in paths.values():
if path.exists():
path.unlink()
split_for = split_groups(records, args.seed)
counts = Counter()
databases = {split: connect(path) for split, path in paths.items()}
for oxide, polymorph, calculation in records:
split = split_for[(oxide, polymorph)]
databases[split].write(
to_atoms(oxide, polymorph, calculation),
oxide=oxide,
polymorph=polymorph,
xc="PBE",
)
counts[split] += 1
manifest = {
"source": str(args.input.resolve()),
"seed": args.seed,
"split_strategy": "oxide-polymorph grouped 80/10/10",
"stress_source_unit": "GPa",
"stress_output_unit": "eV/Angstrom^3",
"counts": dict(counts),
"groups": {f"{oxide}/{polymorph}": split for (oxide, polymorph), split in sorted(split_for.items())},
}
with (args.output / "manifest.json").open("w") as handle:
json.dump(manifest, handle, indent=2, sort_keys=True)
print(json.dumps(manifest["counts"], sort_keys=True))
print(f"saved dataset: {args.output.resolve()}")
if __name__ == "__main__":
main()
|