73 lines
2.3 KiB
Python
Executable File
73 lines
2.3 KiB
Python
Executable File
"""
|
||
Property Prediction Module
|
||
|
||
This module provides functions for predicting properties of crystal structures.
|
||
"""
|
||
|
||
import asyncio
|
||
import torch
|
||
import numpy as np
|
||
from ase.units import GPa
|
||
from mattersim.forcefield import MatterSimCalculator
|
||
|
||
from mars_toolkit.core.llm_tools import llm_tool
|
||
from mars_toolkit.compute.structure_opt import convert_structure
|
||
|
||
@llm_tool(
|
||
name="predict_properties",
|
||
description="Predict energy, forces, and stress of crystal structures based on CIF string",
|
||
)
|
||
async def predict_properties(cif_content: str) -> str:
|
||
"""
|
||
Use MatterSim to predict energy, forces, and stress of crystal structures.
|
||
|
||
Args:
|
||
cif_content: Crystal structure string in CIF format
|
||
|
||
Returns:
|
||
String containing prediction results
|
||
"""
|
||
# 使用asyncio.to_thread异步执行可能阻塞的操作
|
||
def run_prediction():
|
||
# 使用 convert_structure 函数将 CIF 字符串转换为 Atoms 对象
|
||
structure = convert_structure("cif", cif_content)
|
||
if structure is None:
|
||
return "Unable to parse CIF string. Please check if the format is correct."
|
||
|
||
# 设置设备
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
||
# 使用 MatterSimCalculator 计算属性
|
||
structure.calc = MatterSimCalculator(device=device)
|
||
|
||
# 直接获取能量、力和应力
|
||
energy = structure.get_potential_energy()
|
||
forces = structure.get_forces()
|
||
stresses = structure.get_stress(voigt=False)
|
||
|
||
# 计算每原子能量
|
||
num_atoms = len(structure)
|
||
energy_per_atom = energy / num_atoms
|
||
|
||
# 计算应力（GPa和eV/A^3格式）
|
||
stresses_ev_a3 = stresses
|
||
stresses_gpa = stresses / GPa
|
||
|
||
# 构建返回的提示信息
|
||
prompt = f"""
|
||
## {structure.get_chemical_formula()} Crystal Structure Property Prediction Results
|
||
|
||
Prediction results using the provided CIF structure:
|
||
|
||
- Total Energy (eV): {energy}
|
||
- Energy per Atom (eV/atom): {energy_per_atom:.4f}
|
||
- Forces (eV/Angstrom): {forces[0]} # Forces on the first atom
|
||
- Stress (GPa): {stresses_gpa[0][0]} # First component of the stress tensor
|
||
- Stress (eV/A^3): {stresses_ev_a3[0][0]} # First component of the stress tensor
|
||
|
||
"""
|
||
return prompt
|
||
|
||
# 异步执行预测操作
|
||
return await asyncio.to_thread(run_prediction)
|