File size: 4,795 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 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 | """Compute the ASE energy-volume curve from the official NequIP example."""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import numpy as np
import torch
from ase.build import bulk
from onescience.utils.nequip.integrations.ase import NequIPCalculator
def default_compiled_model() -> str | None:
models_dir = os.environ.get("ONESCIENCE_MODELS_DIR")
if not models_dir:
return None
return str(Path(models_dir) / "NequIP" / "NequIP-OAM-L-0.1.nequip.pth")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--compiled-model", default=default_compiled_model())
parser.add_argument("--device", default="cuda")
parser.add_argument("--element", default="Si")
parser.add_argument("--crystal-structure", default="diamond")
parser.add_argument("--lattice-constant", type=float, default=5.43)
parser.add_argument("--supercell", type=int, default=3)
parser.add_argument("--scale-min", type=float, default=0.95)
parser.add_argument("--scale-max", type=float, default=1.05)
parser.add_argument("--num-points", type=int, default=10)
parser.add_argument("--output", default="outputs/energy_volume.json")
parser.add_argument("--plot", default="outputs/energy_volume.png")
args = parser.parse_args()
if not args.compiled_model:
parser.error("--compiled-model is required when ONESCIENCE_MODELS_DIR is unset")
compiled_model = Path(args.compiled_model).expanduser().resolve()
if not compiled_model.is_file():
parser.error(f"compiled model not found: {compiled_model}")
if args.num_points < 2:
parser.error("--num-points must be at least 2")
if args.supercell < 1:
parser.error("--supercell must be positive")
calculator = NequIPCalculator.from_compiled_model(
compile_path=str(compiled_model),
chemical_species_to_atom_type_map={args.element: args.element},
device=args.device,
)
points = []
for scale in np.linspace(args.scale_min, args.scale_max, args.num_points):
atoms = bulk(
args.element,
crystalstructure=args.crystal_structure,
a=args.lattice_constant * float(scale),
cubic=True,
)
atoms *= (args.supercell,) * 3
atoms.calc = calculator
energy = float(atoms.get_potential_energy())
forces = atoms.get_forces()
points.append(
{
"scale": float(scale),
"volume_angstrom3": float(atoms.get_volume()),
"energy_ev": energy,
"energy_ev_per_atom": energy / len(atoms),
"max_force_ev_per_angstrom": float(
np.linalg.norm(forces, axis=1).max()
),
}
)
energies = np.asarray([point["energy_ev"] for point in points])
volumes = np.asarray([point["volume_angstrom3"] for point in points])
minimum_index = int(np.argmin(energies))
result = {
"compiled_model": str(compiled_model),
"device": args.device,
"device_name": torch.cuda.get_device_name(0)
if args.device.startswith("cuda") and torch.cuda.is_available()
else "cpu",
"element": args.element,
"crystal_structure": args.crystal_structure,
"base_lattice_constant_angstrom": args.lattice_constant,
"supercell": [args.supercell] * 3,
"num_atoms": len(atoms),
"points": points,
"sampled_minimum": points[minimum_index],
}
output_path = Path(args.output).expanduser().resolve()
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
if args.plot:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
plot_path = Path(args.plot).expanduser().resolve()
plot_path.parent.mkdir(parents=True, exist_ok=True)
plt.figure(figsize=(8, 6))
plt.plot(volumes, energies, marker="o", label="E-V Curve")
plt.xlabel("Volume (Angstrom^3)", fontsize=14)
plt.ylabel("Energy (eV)", fontsize=14)
plt.title(f"Energy-Volume Curve for Cubic {args.element}", fontsize=16)
plt.legend(fontsize=12)
plt.grid()
plt.tight_layout()
plt.savefig(plot_path, dpi=160)
plt.close()
result["plot"] = str(plot_path)
output_path.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print(f"points: {len(points)}")
print(f"atoms per point: {result['num_atoms']}")
print(f"sampled minimum: {result['sampled_minimum']}")
print(f"result: {output_path}")
if __name__ == "__main__":
main()
|