File size: 1,126 Bytes
45e5a32 | 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 | import sys
from pathlib import Path
# 把本仓库根目录放到 sys.path 最前面，避免 PYTHONPATH 中其他同名 model 包干扰
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import model # noqa: E402
import ase
from ase.build import bulk
import torch
from pymatgen.core.structure import Structure
from onescience.utils.matris import StructOptimizer
model_name = "matris_10m_oam"
device = "cuda" if torch.cuda.is_available() else "cpu"
matris_opt = StructOptimizer(
model=model_name,
task = "efsm",
optimizer = "FIRE", # FIRE, BFGS ...
device=device
)
atom = Structure.from_file("cif_file/demo.cif")
max_steps = 500
fmax = 0.05
opt_result = matris_opt.relax(
atoms=atom, # pymatgen.Structure or ase.Atoms
verbose=True,
steps=max_steps,
fmax=fmax,
relax_cell=max_steps > 0,
ase_filter="FrechetCellFilter",
)
trajectory = opt_result['trajectory']
energy = trajectory.energies[-1]
force = trajectory.forces[-1]
stress = trajectory.stresses[-1]
magmom = trajectory.magmoms[-1]
final_structure = opt_result['final_structure']
|