| """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)) |
| |
| 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() |
|
|
