初次提交

sci_mcp/material_mcp/mattersim_pred/property_pred_tools.py

@@ -0,0 +1,73 @@
+"""
+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 ...core.llm_tools import llm_tool
+from ..support.utils import convert_structure,read_structure_from_file_name_or_content_string
+
+@llm_tool(
+    name="predict_properties_MatterSim",
+    description="Predict energy, forces, and stress of crystal structures using MatterSim model based on CIF string",
+)
+async def predict_properties_MatterSim(structure_source: str) -> str:
+    """
+    Use MatterSim model to predict energy, forces, and stress of crystal structures.
+
+    Args:
+        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string
+
+    Returns:
+        String containing prediction results
+    """
+    # 使用asyncio.to_thread异步执行可能阻塞的操作
+    def run_prediction():
+        # 使用 convert_structure 函数将 CIF 字符串转换为 Atoms 对象
+        structure_content,content_format=read_structure_from_file_name_or_content_string(structure_source)
+        structure = convert_structure(content_format,  structure_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)