File size: 1,037 Bytes
3e02ab8 | 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 | """Generate a tiny extxyz smoke dataset for NequIP demo training."""
from __future__ import annotations
import os
from pathlib import Path
import numpy as np
from ase import Atoms
from ase.build import bulk
from ase.io import write
def main() -> None:
out_dir = Path(__file__).parent / "reference_data"
out_dir.mkdir(parents=True, exist_ok=True)
out_file = out_dir / "smoke.xyz"
rng = np.random.default_rng(123)
structures = []
# A few small Cu clusters with random displacements.
base = bulk("Cu", "fcc", a=3.6) * (2, 2, 2)
for i in range(8):
atoms = base.copy()
atoms.positions += rng.normal(scale=0.05, size=atoms.positions.shape)
atoms.info["energy"] = float(-len(atoms) * 3.5 + rng.normal(scale=0.5))
atoms.arrays["forces"] = rng.normal(scale=0.1, size=atoms.positions.shape)
structures.append(atoms)
write(out_file, structures, format="extxyz")
print(f"Wrote {len(structures)} structures to {out_file}")
if __name__ == "__main__":
main()
|