diff --git a/src/abacusagent/modules/structure_generator.py b/src/abacusagent/modules/structure_generator.py index cd5eb6c..e6b3ee4 100644 --- a/src/abacusagent/modules/structure_generator.py +++ b/src/abacusagent/modules/structure_generator.py @@ -6,16 +6,18 @@ from abacusagent.modules.submodules.structure_generator import generate_molecule_structure as _generate_molecule_structure from abacusagent.modules.submodules.structure_generator import generate_bulk_structure_from_wyckoff_position as _generate_bulk_structure_from_wyckoff_position from abacusagent.modules.submodules.structure_generator import get_ieee_standard_structure as _get_ieee_standard_structure +from abacusagent.modules.submodules.structure_generator import materials_project_download as _materials_project_download + @mcp.tool() def generate_bulk_structure(element: str, - crystal_structure:Literal["sc", "fcc", "bcc","hcp","diamond", "zincblende", "rocksalt"]='fcc', - a:float =None, - c: float =None, - cubic: bool =False, - orthorhombic: bool =False, - file_format: Literal["cif", "poscar"] = "cif", - ) -> Dict[str, Any]: + crystal_structure:Literal["sc", "fcc", "bcc","hcp","diamond", "zincblende", "rocksalt"]='fcc', + a:float =None, + c: float =None, + cubic: bool =False, + orthorhombic: bool =False, + file_format: Literal["cif", "poscar"] = "cif", + ) -> Dict[str, Any]: """ Generate a bulk crystal structure using ASE's `bulk` function. @@ -155,3 +157,23 @@ def get_ieee_standard_structure( - standard_stru_file: The absolute path to the rotated crystal structure file. """ return _get_ieee_standard_structure(stru_file, stru_type) + +@mcp.tool() +def materials_project_download( + material_id: str, + destination_path: Optional[str] = None, +) -> Dict[str, Any]: + """ + Download structure from Materials Project database by material ID. + + Args: + material_id (str): The Materials Project material ID (e.g., 'mp-12345') + destination_path (str, optional): The path to save the downloaded structure file. + If not provided, a temporary file will be created. + + Returns: + A dictionary containing: + - 'structure_file': Path to the downloaded structure file. + - 'material_id': The material ID used for download. + """ + return _materials_project_download(material_id, destination_path) diff --git a/src/abacusagent/modules/submodules/abacus.py b/src/abacusagent/modules/submodules/abacus.py index 851843f..0448aad 100644 --- a/src/abacusagent/modules/submodules/abacus.py +++ b/src/abacusagent/modules/submodules/abacus.py @@ -12,6 +12,9 @@ from abacusagent.init_mcp import mcp from abacusagent.modules.util.comm import generate_work_path, run_abacus, collect_metrics +import tempfile +from pymatgen.ext.matproj import MPRester +from pymatgen.core import Structure def abacus_prepare( stru_file: Path, @@ -614,3 +617,4 @@ def read_abacus_stru(abacus_input_dir: Path): 'move': stru.get_move()} except Exception as e: return {'message': f"Read ABACUS STRU file failed: {e}"} + diff --git a/src/abacusagent/modules/submodules/structure_generator.py b/src/abacusagent/modules/submodules/structure_generator.py index 6c1a066..8b8d6dc 100644 --- a/src/abacusagent/modules/submodules/structure_generator.py +++ b/src/abacusagent/modules/submodules/structure_generator.py @@ -1,91 +1,71 @@ +import os +from pathlib import Path +from typing import Literal, Optional, Dict, Any, List, Tuple, Union + from ase import Atoms from ase.io import read, write from ase.build import molecule from ase.data import chemical_symbols from ase.collections import g2 from pymatgen.core import Structure, Lattice - -from pathlib import Path -from typing import Literal, Optional, Dict, Any, List, Tuple, Union +from pymatgen.ext.matproj import MPRester from abacusagent.modules.util.comm import generate_work_path -# From Introduction to Solid State Physics, 8th edition, by Charles Kittel -ELEMENT_CRYSTAL_STRUCTURES = { - "Li": {"crystal": "bcc", "a": 3.51}, - "Be": {"crystal": "hcp", "a": 2.27, "c": 3.59}, - "Na": {"crystal": "bcc", "a": 4.23}, - "Mg": {"crystal": "hcp", "a": 3.21, "c": 5.21}, - "Al": {"crystal": "fcc", "a": 4.05}, - "Si": {"crystal": "diamond", "a": 5.43}, - "K": {"crystal": "bcc", "a": 5.23}, - "Ca": {"crystal": "fcc", "a": 5.58}, - "Sc": {"crystal": "hcp", "a": 3.31, "c": 5.27}, - "Ti": {"crystal": "hcp", "a": 2.95, "c": 4.68}, - "V": {"crystal": "bcc", "a": 3.03}, - "Cr": {"crystal": "bcc", "a": 2.88}, - "Mn": {"crystal": "bcc", "a": 2.91}, # To be checked - "Fe": {"crystal": "bcc", "a": 2.87}, - "Co": {"crystal": "hcp", "a": 2.51, "c": 4.07}, - "Ni": {"crystal": "fcc", "a": 3.52}, - "Cu": {"crystal": "fcc", "a": 3.61}, - "Zn": {"crystal": "hcp", "a": 2.66, "c": 4.95}, - "Ga": {"crystal": "fcc", "a": 4.50}, # To be checked - "Ge": {"crystal": "diamond", "a": 5.69}, - "Rb": {"crystal": "bcc", "a": 5.58}, - "Sr": {"crystal": "fcc", "a": 6.08}, - "Y": {"crystal": "hcp", "a": 3.65, "c": 5.73}, - "Zr": {"crystal": "hcp", "a": 3.23, "c": 5.15}, - "Nb": {"crystal": "bcc", "a": 3.30}, - "Mo": {"crystal": "bcc", "a": 3.15}, - "Tc": {"crystal": "hcp", "a": 2.74, "c": 4.44}, - "Ru": {"crystal": "hcp", "a": 2.71, "c": 4.28}, - "Rh": {"crystal": "fcc", "a": 3.80}, - "Pd": {"crystal": "fcc", "a": 3.89}, - "Ag": {"crystal": "fcc", "a": 4.09}, - "Cd": {"crystal": "hcp", "a": 2.98, "c": 5.62}, - "In": {"crystal": "bcc", "a": 3.25}, # To be checked - "Sn": {"crystal": "diamond", "a": 6.49}, - "Cs": {"crystal": "bcc", "a": 6.05}, - "Ba": {"crystal": "bcc", "a": 5.02}, - "La": {"crystal": "hcp", "a": 3.75, "c": 5.75}, # To be checked - "Ce": {"crystal": "fcc", "a": 5.16}, - "Pr": {"crystal": "hcp", "a": 3.65, "c": 5.75}, # To be checked - "Nd": {"crystal": "hcp", "a": 3.65, "c": 5.73}, # To be checked - "Pm": {"crystal": "hcp", "a": 3.65, "c": 5.73}, # To be checked - "Sm": {"crystal": "hcp", "a": 3.65, "c": 5.73}, # To be checked - "Eu": {"crystal": "bcc", "a": 4.58}, - "Gd": {"crystal": "hcp", "a": 3.63, "c": 5.78}, - "Tb": {"crystal": "hcp", "a": 3.60, "c": 5.70}, - "Dy": {"crystal": "hcp", "a": 3.59, "c": 5.65}, - "Ho": {"crystal": "hcp", "a": 3.58, "c": 5.62}, - "Er": {"crystal": "hcp", "a": 3.56, "c": 5.59}, - "Tm": {"crystal": "hcp", "a": 3.54, "c": 5.56}, - "Yb": {"crystal": "fcc", "a": 5.48}, - "Lu": {"crystal": "hcp", "a": 3.50, "c": 5.55}, - "Hf": {"crystal": "hcp", "a": 3.19, "c": 5.05}, - "Ta": {"crystal": "bcc", "a": 3.30}, - "W": {"crystal": "bcc", "a": 3.16}, - "Re": {"crystal": "hcp", "a": 2.76, "c": 4.46}, - "Os": {"crystal": "hcp", "a": 2.74, "c": 4.32}, - "Ir": {"crystal": "fcc", "a": 3.84}, - "Pt": {"crystal": "fcc", "a": 3.92}, - "Au": {"crystal": "fcc", "a": 4.08}, - "Hg": {"crystal": "hcp", "a": 2.99, "c": 5.01}, # To be checked - "Tl": {"crystal": "hcp", "a": 3.46, "c": 5.52}, - "Pb": {"crystal": "fcc", "a": 4.95}, -} +def materials_project_download( + material_id: str, + destination_path: Optional[str] = None, +) -> Dict[str, Any]: + """ + Download structure from Materials Project database by material ID. + + Args: + material_id (str): The Materials Project material ID (e.g., 'mp-12345') + destination_path (str, optional): The path to save the downloaded structure file. + If not provided, a temporary file will be created. + + Returns: + A dictionary containing: + - 'structure_file': Path to the downloaded structure file. + - 'material_id': The material ID used for download. + """ + try: + # Get the Materials Project API key from environment variables + api_key = os.environ.get("MP_API_KEY") + if not api_key: + raise ValueError("Materials Project API key not found. Please set MP_API_KEY environment variable.") + + # Connect to Materials Project database + with MPRester(api_key) as mpr: + # Retrieve the structure + structure = mpr.get_structure_by_material_id(material_id) + + # Determine destination path + if destination_path is None: + destination_path = f"{material_id}.cif" + + # Save structure + structure.to(filename=destination_path, fmt="cif") + + return { + "structure_file": destination_path, + "material_id": material_id, + } + + except Exception as e: + import traceback + traceback.print_exc() + return {"message": f"Failed to download structure from Materials Project: {e}"} -#@mcp.tool() def generate_bulk_structure(element: str, - crystal_structure:Literal["sc", "fcc", "bcc","hcp","diamond", "zincblende", "rocksalt"]='fcc', - a:float =None, - c: float =None, - cubic: bool =False, - orthorhombic: bool =False, - file_format: Literal["cif", "poscar"] = "cif", - ) -> Dict[str, Any]: + crystal_structure:Literal["sc", "fcc", "bcc","hcp","diamond", "zincblende", "rocksalt"]='fcc', + a:float =None, + c: float =None, + cubic: bool =False, + orthorhombic: bool =False, + file_format: Literal["cif", "poscar"] = "cif", + ) -> Dict[str, Any]: """ Generate a bulk crystal structure using ASE's `bulk` function. @@ -163,7 +143,6 @@ def generate_bulk_structure(element: str, except Exception as e: return {"message": f"Generating bulk structure failed: {e}"} -#@mcp.tool() def generate_bulk_structure_from_wyckoff_position( a: float, b: float, @@ -213,7 +192,6 @@ def generate_bulk_structure_from_wyckoff_position( except Exception as e: return {"message": f"Generating bulk structure from Wyckoff position failed: {e}"} -#@mcp.tool() def generate_molecule_structure( molecule_name: Literal['PH3', 'P2', 'CH3CHO', 'H2COH', 'CS', 'OCHCHO', 'C3H9C', 'CH3COF', 'CH3CH2OCH3', 'HCOOH', 'HCCl3', 'HOCl', 'H2', 'SH2', 'C2H2', 'C4H4NH', diff --git a/tests/integrate_test/test_materials_project.py b/tests/integrate_test/test_materials_project.py new file mode 100644 index 0000000..d83fbac --- /dev/null +++ b/tests/integrate_test/test_materials_project.py @@ -0,0 +1,51 @@ +import os +import tempfile +import unittest +from pathlib import Path +from abacusagent.modules.submodules.structure_generator import materials_project_download + +class TestMaterialsProjectDownload(unittest.TestCase): + + def setUp(self): + self.api_key = os.environ.get('MP_API_KEY', 'test_key') + os.environ['MP_API_KEY'] = self.api_key + + def test_materials_project_download_basic(self): + result = materials_project_download(material_id="mp-1234") + + self.assertIn('structure_file', result) + self.assertIn('material_id', result) + self.assertEqual(result['material_id'], "mp-1234") + + # 清理下载的文件 + if 'structure_file' in result and os.path.exists(result['structure_file']): + os.unlink(result['structure_file']) + + def test_materials_project_download_with_custom_path(self): + """测试指定输出路径的功能""" + with tempfile.NamedTemporaryFile(suffix='.cif', delete=False) as tmp: + temp_path = tmp.name + + try: + result = materials_project_download( + material_id="mp-1234", + destination_path=temp_path + ) + + self.assertEqual(result['structure_file'], temp_path) + self.assertTrue(os.path.exists(temp_path)) + finally: + if os.path.exists(temp_path): + os.unlink(temp_path) + + def test_materials_project_download_error_handling(self): + result = materials_project_download(material_id="mp-999999999") + + self.assertIn('message', result) + + # 如果下载成功（虽然不应该），清理文件 + if 'structure_file' in result and os.path.exists(result.get('structure_file', '')): + os.unlink(result['structure_file']) + +if __name__ == '__main__': + unittest.main()