初次提交

2025-05-09 14:16:33 +08:00
commit 3a50afeec4
56 changed files with 9224 additions and 0 deletions
--- a/test_tools/material/test_structure_opt.py
+++ b/test_tools/material/test_structure_opt.py
@@ -0,0 +1,118 @@
+"""
+测试晶体结构优化工具函数
+
+此脚本测试改进后的optimize_crystal_structure函数，
+该函数接受单一的file_name_or_content_string参数，可以是文件路径或直接的结构内容。
+"""
+
+import sys
+import asyncio
+import os
+import tempfile
+sys.path.append("/home/ubuntu/sas0/lzy/multi_mcp_server")
+
+from sci_mcp.material_mcp.fairchem_structure_opt.structure_opt_tools import optimize_crystal_structure
+from sci_mcp.core.config import material_config
+
+# 简单的CIF结构示例
+SAMPLE_CIF = """
+data_SrTiO3
+_cell_length_a 3.905
+_cell_length_b 3.905
+_cell_length_c 3.905
+_cell_angle_alpha 90
+_cell_angle_beta 90
+_cell_angle_gamma 90
+_symmetry_space_group_name_H-M 'P m -3 m'
+_symmetry_Int_Tables_number 221
+loop_
+_atom_site_label
+_atom_site_type_symbol
+_atom_site_fract_x
+_atom_site_fract_y
+_atom_site_fract_z
+Sr1 Sr 0.0 0.0 0.0
+Ti1 Ti 0.5 0.5 0.5
+O1 O 0.5 0.5 0.0
+O2 O 0.5 0.0 0.5
+O3 O 0.0 0.5 0.5
+"""
+
+async def test_with_content():
+    """测试使用直接结构内容"""
+    print("\n=== 测试使用直接结构内容 ===")
+    result = await optimize_crystal_structure(
+        file_name_or_content_string=SAMPLE_CIF,
+        format_type="cif",
+        optimization_level="quick"
+    )
+    print(result)
+
+async def test_with_file():
+    """测试使用文件路径（如果文件存在）"""
+    print("\n=== 测试使用文件路径 ===")
+    # 创建临时CIF文件
+
+    with tempfile.NamedTemporaryFile(suffix=".cif", mode="w", delete=False) as tmp_file:
+        tmp_file.write(SAMPLE_CIF)
+        tmp_path = tmp_file.name
+
+    try:
+        result = await optimize_crystal_structure(
+            file_name_or_content_string=tmp_path,
+            format_type="auto",
+            optimization_level="quick"
+        )
+        print(result)
+    finally:
+        # 清理临时文件
+        if os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+
+async def test_with_temp_file():
+    """测试使用临时目录中的文件名"""
+    print("\n=== 测试使用临时目录中的文件名 ===")
+
+    # 确保临时目录存在
+    os.makedirs(material_config.TEMP_ROOT, exist_ok=True)
+
+    # 在临时目录中创建文件
+    temp_filename = "test_structure.cif"
+    temp_filepath = os.path.join(material_config.TEMP_ROOT, temp_filename)
+
+    with open(temp_filepath, 'w', encoding='utf-8') as f:
+        f.write(SAMPLE_CIF)
+
+    try:
+        # 只传递文件名，而不是完整路径
+        result = await optimize_crystal_structure(
+            file_name_or_content_string=temp_filename,
+            format_type="auto",
+            optimization_level="quick"
+        )
+        print(result)
+    finally:
+        # 清理临时文件
+        if os.path.exists(temp_filepath):
+            os.unlink(temp_filepath)
+
+async def test_auto_format():
+    """测试自动格式检测"""
+    print("\n=== 测试自动格式检测 ===")
+    result = await optimize_crystal_structure(
+        file_name_or_content_string=SAMPLE_CIF,
+        format_type="auto"
+    )
+    print(result)
+
+async def main():
+    """运行所有测试"""
+    print("测试改进后的optimize_crystal_structure函数")
+
+    await test_with_content()
+    await test_with_file()
+    await test_with_temp_file()
+    await test_auto_format()
+
+if __name__ == "__main__":
+    asyncio.run(main())