119 lines
3.2 KiB
Python
119 lines
3.2 KiB
Python
"""
|
||
测试晶体结构优化工具函数
|
||
|
||
此脚本测试改进后的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())
|