| """Fit an elastic tensor from Equiformer V3 stress predictions.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| from pathlib import Path |
|
|
| os.environ.setdefault( |
| "ONESCIENCE_EQUIFORMER_V3_JD_PATH", |
| str(Path(__file__).resolve().parent / "weight" / "Jd.pt"), |
| ) |
|
|
| from ase.build import bulk |
| from ase.io import read, write |
|
|
| from onescience.utils.equiformer_v3 import ( |
| EquiformerV3Calculator, |
| relax_structure, |
| run_elastic_workflow, |
| write_workflow_result, |
| ) |
|
|
|
|
| def default_checkpoint() -> str: |
| return str( |
| Path(__file__).resolve().parent |
| / "weight" |
| / "omat24-mptrj-salex_gradient.pt" |
| ) |
|
|
|
|
| def load_structure(path: str | None): |
| if path: |
| return read(path) |
| return bulk("Cu") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--checkpoint") |
| parser.add_argument( |
| "--input", |
| help="periodic structure readable by ASE; defaults to built-in bulk Cu", |
| ) |
| parser.add_argument( |
| "--relax", |
| action="store_true", |
| help="relax the cell and atomic positions before applying strains", |
| ) |
| parser.add_argument( |
| "--normal-strains", type=float, nargs="+", default=(-0.01, 0.01) |
| ) |
| parser.add_argument( |
| "--shear-strains", type=float, nargs="+", default=(-0.02, 0.02) |
| ) |
| parser.add_argument( |
| "--relax-positions", |
| action="store_true", |
| help="relax atomic positions at fixed cell for every deformation", |
| ) |
| parser.add_argument("--relax-fmax", type=float, default=0.02) |
| parser.add_argument("--relax-steps", type=int, default=100) |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--output", default="outputs/elastic.json") |
| args = parser.parse_args() |
|
|
| checkpoint = args.checkpoint or default_checkpoint() |
| calculator = EquiformerV3Calculator.from_checkpoint( |
| checkpoint, device=args.device |
| ) |
| structure = load_structure(args.input) |
| initial_relaxation = None |
| if args.relax: |
| structure, initial_relaxation = relax_structure( |
| structure, |
| calculator, |
| relax_cell=True, |
| fmax=args.relax_fmax, |
| steps=args.relax_steps, |
| ) |
| relaxed_output = Path(args.output).with_name("elastic_relaxed.cif") |
| relaxed_output.parent.mkdir(parents=True, exist_ok=True) |
| write(relaxed_output, structure) |
| result = run_elastic_workflow( |
| structure, |
| calculator, |
| normal_strains=args.normal_strains, |
| shear_strains=args.shear_strains, |
| relax_positions=args.relax_positions, |
| relax_fmax=args.relax_fmax, |
| relax_steps=args.relax_steps, |
| ) |
| result["checkpoint"] = str(checkpoint) |
| result["initial_relaxation"] = initial_relaxation |
| output = write_workflow_result(result, args.output) |
|
|
| print("formula:", result["formula"]) |
| print("bulk modulus, Hill (GPa):", result["bulk_modulus_gpa"]["hill"]) |
| print("result:", output) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
