初次提交

sci_mcp/material_mcp/matgl_tools/__init__.py

@@ -0,0 +1,8 @@
+"""
+MatGL Tools Module
+
+This module provides tools for material structure relaxation and property prediction
+using MatGL (Materials Graph Library) models.
+"""
+
+from .matgl_tools import *

sci_mcp/material_mcp/matgl_tools/matgl_tools.py

@@ -0,0 +1,487 @@
+"""
+MatGL Tools Module
+
+This module provides tools for material structure relaxation and property prediction
+using MatGL (Materials Graph Library) models.
+"""
+from __future__ import annotations
+from ...core.config import material_config
+
+
+import warnings
+import json
+from typing import Dict, List, Union, Optional, Any
+
+import torch
+from pymatgen.core import Lattice, Structure
+from pymatgen.ext.matproj import MPRester
+from pymatgen.io.ase import AseAtomsAdaptor
+
+import matgl
+from matgl.ext.ase import Relaxer, MolecularDynamics, PESCalculator
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+
+from ...core.llm_tools import llm_tool
+import os
+from ..support.utils import read_structure_from_file_name_or_content_string
+# To suppress warnings for clearer output
+warnings.simplefilter("ignore")
+
+
+@llm_tool(name="relax_crystal_structure_M3GNet",
+          description="Optimize crystal structure geometry using M3GNet universal potential from a structure file or content string")
+async def relax_crystal_structure_M3GNet(
+    structure_source: str,
+    fmax: float = 0.01
+) -> str:
+    """
+    Optimize crystal structure geometry to find its equilibrium configuration.
+
+    Uses M3GNet universal potential for fast and accurate structure relaxation without DFT.
+    Accepts a structure file or content string.
+
+    Args:
+        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
+        fmax: Maximum force tolerance for convergence in eV/Å (default: 0.01).
+
+    Returns:
+        A Markdown formatted string with the relaxation results or an error message.
+    """
+    try:
+        # 使用read_structure_from_file_name_or_content_string函数读取结构
+        structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
+
+        # 使用pymatgen读取结构
+        structure = Structure.from_str(structure_content, fmt=content_format)
+
+        if structure is None:
+            return "Error: Failed to obtain a valid structure"
+
+        # Load the M3GNet universal potential model
+        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+
+        # Create a relaxer and relax the structure
+        relaxer = Relaxer(potential=pot)
+        relax_results = relaxer.relax(structure, fmax=fmax)
+
+        # Get the relaxed structure
+        relaxed_structure = relax_results["final_structure"]
+        reduced_formula = relaxed_structure.composition.reduced_formula
+
+        # 添加结构信息
+        lattice_info = relaxed_structure.lattice
+        volume = relaxed_structure.volume
+        density = relaxed_structure.density
+        symmetry = relaxed_structure.get_space_group_info()
+
+        # 构建原子位置表格
+        sites_table = "  #  SP           a         b         c\n"
+        sites_table += "---  ----  --------  --------  --------\n"
+        for i, site in enumerate(relaxed_structure):
+            frac_coords = site.frac_coords
+            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"
+
+        return (f"## Structure Relaxation\n\n"
+                f"- **Structure**: `{reduced_formula}`\n"
+                f"- **Force Tolerance**: `{fmax} eV/Å`\n"
+                f"- **Status**: `Successfully relaxed`\n\n"
+                f"### Relaxed Structure Information\n\n"
+                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
+                f"- **Volume**: `{volume:.2f} ų`\n"
+                f"- **Density**: `{density:.2f} g/cm³`\n"
+                f"- **Lattice Parameters**:\n"
+                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
+                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
+                f"### Atomic Positions (Fractional Coordinates)\n\n"
+                f"```\n"
+                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
+                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
+                f"pbc   :       {relaxed_structure.lattice.pbc[0]!s:5s}       {relaxed_structure.lattice.pbc[1]!s:5s}       {relaxed_structure.lattice.pbc[2]!s:5s}\n"
+                f"Sites ({len(relaxed_structure)})\n"
+                f"{sites_table}```\n")
+    except Exception as e:
+        return f"Error during structure relaxation: {str(e)}"
+
+
+# 内部函数，用于结构优化，返回结构对象而不是格式化字符串
+async def _relax_crystal_structure_M3GNet_internal(
+    structure_source: str,
+    fmax: float = 0.01
+) -> Union[Structure, str]:
+    """
+    内部使用的结构优化函数，返回结构对象而不是格式化字符串。
+
+    Args:
+        structure_source: 结构文件名或内容字符串
+        fmax: 力收敛阈值 (eV/Å)
+
+    Returns:
+        优化后的结构对象或错误信息
+    """
+    try:
+        # 使用read_structure_from_file_name_or_content_string函数读取结构
+        structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
+
+        # 使用pymatgen读取结构
+        structure = Structure.from_str(structure_content, fmt=content_format)
+
+        if structure is None:
+            return "Error: Failed to obtain a valid structure"
+
+        # Load the M3GNet universal potential model
+        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+
+        # Create a relaxer and relax the structure
+        relaxer = Relaxer(potential=pot)
+        relax_results = relaxer.relax(structure, fmax=fmax)
+
+        # Get the relaxed structure
+        relaxed_structure = relax_results["final_structure"]
+
+        return relaxed_structure
+    except Exception as e:
+        return f"Error during structure relaxation: {str(e)}"
+
+
+@llm_tool(name="predict_formation_energy_M3GNet",
+          description="Predict the formation energy of a crystal structure using the M3GNet formation energy model from a structure file or content string, with optional structure optimization")
+async def predict_formation_energy_M3GNet(
+    structure_source: str,
+    optimize_structure: bool = True,
+    fmax: float = 0.01
+) -> str:
+    """
+    Predict the formation energy of a crystal structure using the M3GNet formation energy model.
+
+    Args:
+        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
+        optimize_structure: Whether to optimize the structure before prediction (default: True).
+        fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).
+
+    Returns:
+        A Markdown formatted string containing the predicted formation energy in eV/atom or an error message.
+    """
+    try:
+        # 获取结构（优化或不优化）
+        if optimize_structure:
+            # 使用内部函数优化结构
+            structure = await _relax_crystal_structure_M3GNet_internal(
+                structure_source=structure_source,
+                fmax=fmax
+            )
+
+            # 检查优化是否成功
+            if isinstance(structure, str) and structure.startswith("Error"):
+                return structure
+        else:
+            # 直接读取结构，不进行优化
+            structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
+            structure = Structure.from_str(structure_content, fmt=content_format)
+
+        if structure is None:
+            return "Error: Failed to obtain a valid structure"
+
+        # 加载预训练模型
+        model = matgl.load_model("M3GNet-MP-2018.6.1-Eform")
+
+        # 预测形成能
+        eform = model.predict_structure(structure)
+        reduced_formula = structure.composition.reduced_formula
+
+        # 构建结果字符串
+        optimization_status = "optimized" if optimize_structure else "non-optimized"
+
+        # 添加结构信息
+        lattice_info = structure.lattice
+        volume = structure.volume
+        density = structure.density
+        symmetry = structure.get_space_group_info()
+
+        # 构建原子位置表格
+        sites_table = "  #  SP           a         b         c\n"
+        sites_table += "---  ----  --------  --------  --------\n"
+        for i, site in enumerate(structure):
+            frac_coords = site.frac_coords
+            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"
+
+        return (f"## Formation Energy Prediction\n\n"
+                f"- **Structure**: `{reduced_formula}`\n"
+                f"- **Structure Status**: `{optimization_status}`\n"
+                f"- **Formation Energy**: `{float(eform):.3f} eV/atom`\n\n"
+                f"### Structure Information\n\n"
+                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
+                f"- **Volume**: `{volume:.2f} ų`\n"
+                f"- **Density**: `{density:.2f} g/cm³`\n"
+                f"- **Lattice Parameters**:\n"
+                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
+                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
+                f"### Atomic Positions (Fractional Coordinates)\n\n"
+                f"```\n"
+                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
+                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
+                f"pbc   :       {structure.lattice.pbc[0]!s:5s}       {structure.lattice.pbc[1]!s:5s}       {structure.lattice.pbc[2]!s:5s}\n"
+                f"Sites ({len(structure)})\n"
+                f"{sites_table}```\n")
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+@llm_tool(name="run_molecular_dynamics_M3GNet",
+          description="Run molecular dynamics simulation on a crystal structure using M3GNet universal potential, with optional structure optimization")
+async def run_molecular_dynamics_M3GNet(
+    structure_source: str,
+    temperature_K: float = 300,
+    steps: int = 100,
+    optimize_structure: bool = True,
+    fmax: float = 0.01
+) -> str:
+    """
+    Run molecular dynamics simulation on a crystal structure using M3GNet universal potential.
+
+    Args:
+        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
+        temperature_K: Temperature for MD simulation in Kelvin (default: 300).
+        steps: Number of MD steps to run (default: 100).
+        optimize_structure: Whether to optimize the structure before simulation (default: True).
+        fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).
+
+    Returns:
+        A Markdown formatted string containing the simulation results, including final potential energy.
+    """
+    try:
+        # 获取结构（优化或不优化）
+        if optimize_structure:
+            # 使用内部函数优化结构
+            structure = await _relax_crystal_structure_M3GNet_internal(
+                structure_source=structure_source,
+                fmax=fmax
+            )
+
+            # 检查优化是否成功
+            if isinstance(structure, str) and structure.startswith("Error"):
+                return structure
+        else:
+            # 直接读取结构，不进行优化
+            structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
+            structure = Structure.from_str(structure_content, fmt=content_format)
+
+        if structure is None:
+            return "Error: Failed to obtain a valid structure"
+
+        # Load the M3GNet universal potential model
+        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+
+        # Convert pymatgen structure to ASE atoms
+        ase_adaptor = AseAtomsAdaptor()
+        atoms = ase_adaptor.get_atoms(structure)
+
+        # Initialize the velocity according to Maxwell Boltzmann distribution
+        MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K)
+
+        # Create the MD class and run simulation
+        driver = MolecularDynamics(atoms, potential=pot, temperature=temperature_K)
+        driver.run(steps)
+
+        # Get final potential energy
+        final_energy = atoms.get_potential_energy()
+
+        # Get final structure
+        final_structure = ase_adaptor.get_structure(atoms)
+        reduced_formula = final_structure.composition.reduced_formula
+
+        # 构建结果字符串
+        optimization_status = "optimized" if optimize_structure else "non-optimized"
+
+        # 添加结构信息
+        lattice_info = final_structure.lattice
+        volume = final_structure.volume
+        density = final_structure.density
+        symmetry = final_structure.get_space_group_info()
+
+        # 构建原子位置表格
+        sites_table = "  #  SP           a         b         c\n"
+        sites_table += "---  ----  --------  --------  --------\n"
+        for i, site in enumerate(final_structure):
+            frac_coords = site.frac_coords
+            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"
+
+        return (f"## Molecular Dynamics Simulation\n\n"
+                f"- **Structure**: `{reduced_formula}`\n"
+                f"- **Structure Status**: `{optimization_status}`\n"
+                f"- **Temperature**: `{temperature_K} K`\n"
+                f"- **Steps**: `{steps}`\n"
+                f"- **Final Potential Energy**: `{float(final_energy):.3f} eV`\n\n"
+                f"### Final Structure Information\n\n"
+                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
+                f"- **Volume**: `{volume:.2f} ų`\n"
+                f"- **Density**: `{density:.2f} g/cm³`\n"
+                f"- **Lattice Parameters**:\n"
+                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
+                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
+                f"### Atomic Positions (Fractional Coordinates)\n\n"
+                f"```\n"
+                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
+                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
+                f"pbc   :       {final_structure.lattice.pbc[0]!s:5s}       {final_structure.lattice.pbc[1]!s:5s}       {final_structure.lattice.pbc[2]!s:5s}\n"
+                f"Sites ({len(final_structure)})\n"
+                f"{sites_table}```\n")
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+@llm_tool(name="calculate_single_point_energy_M3GNet",
+          description="Calculate single point energy of a crystal structure using M3GNet universal potential, with optional structure optimization")
+async def calculate_single_point_energy_M3GNet(
+    structure_source: str,
+    optimize_structure: bool = True,
+    fmax: float = 0.01
+) -> str:
+    """
+    Calculate single point energy of a crystal structure using M3GNet universal potential.
+
+    Args:
+        structure_source: The name of the structure file (e.g., POSCAR, CIF) or the content string.
+        optimize_structure: Whether to optimize the structure before calculation (default: True).
+        fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).
+
+    Returns:
+        A Markdown formatted string containing the calculated potential energy in eV.
+    """
+    try:
+        # 获取结构（优化或不优化）
+        if optimize_structure:
+            # 使用内部函数优化结构
+            structure = await _relax_crystal_structure_M3GNet_internal(
+                structure_source=structure_source,
+                fmax=fmax
+            )
+
+            # 检查优化是否成功
+            if isinstance(structure, str) and structure.startswith("Error"):
+                return structure
+        else:
+            # 直接读取结构，不进行优化
+            structure_content, content_format = read_structure_from_file_name_or_content_string(structure_source)
+            structure = Structure.from_str(structure_content, fmt=content_format)
+
+        if structure is None:
+            return "Error: Failed to obtain a valid structure"
+
+        # Load the M3GNet universal potential model
+        pot = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+
+        # Convert pymatgen structure to ASE atoms
+        ase_adaptor = AseAtomsAdaptor()
+        atoms = ase_adaptor.get_atoms(structure)
+
+        # Set up the calculator for atoms object
+        calc = PESCalculator(pot)
+        atoms.set_calculator(calc)
+
+        # Calculate potential energy
+        energy = atoms.get_potential_energy()
+        reduced_formula = structure.composition.reduced_formula
+
+        # 构建结果字符串
+        optimization_status = "optimized" if optimize_structure else "non-optimized"
+
+        # 添加结构信息
+        lattice_info = structure.lattice
+        volume = structure.volume
+        density = structure.density
+        symmetry = structure.get_space_group_info()
+
+        # 构建原子位置表格
+        sites_table = "  #  SP           a         b         c\n"
+        sites_table += "---  ----  --------  --------  --------\n"
+        for i, site in enumerate(structure):
+            frac_coords = site.frac_coords
+            sites_table += f"{i:3d}  {site.species_string:4s}  {frac_coords[0]:8.6f}  {frac_coords[1]:8.6f}  {frac_coords[2]:8.6f}\n"
+
+        return (f"## Single Point Energy Calculation\n\n"
+                f"- **Structure**: `{reduced_formula}`\n"
+                f"- **Structure Status**: `{optimization_status}`\n"
+                f"- **Potential Energy**: `{float(energy):.3f} eV`\n\n"
+                f"### Structure Information\n\n"
+                f"- **Space Group**: `{symmetry[0]} (#{symmetry[1]})`\n"
+                f"- **Volume**: `{volume:.2f} ų`\n"
+                f"- **Density**: `{density:.2f} g/cm³`\n"
+                f"- **Lattice Parameters**:\n"
+                f"  - a = `{lattice_info.a:.6f} Å`, b = `{lattice_info.b:.6f} Å`, c = `{lattice_info.c:.6f} Å`\n"
+                f"  - α = `{lattice_info.alpha:.6f}°`, β = `{lattice_info.beta:.6f}°`, γ = `{lattice_info.gamma:.6f}°`\n\n"
+                f"### Atomic Positions (Fractional Coordinates)\n\n"
+                f"```\n"
+                f"abc   :   {lattice_info.a:.6f}   {lattice_info.b:.6f}   {lattice_info.c:.6f}\n"
+                f"angles:  {lattice_info.alpha:.6f}  {lattice_info.beta:.6f}  {lattice_info.gamma:.6f}\n"
+                f"pbc   :       {structure.lattice.pbc[0]!s:5s}       {structure.lattice.pbc[1]!s:5s}       {structure.lattice.pbc[2]!s:5s}\n"
+                f"Sites ({len(structure)})\n"
+                f"{sites_table}```\n")
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+#Error: Bad serialized model or bad model name. It is possible that you have an older model cached. Please clear your cache by running `python -c "import matgl; matgl.clear_cache()"`
+# @llm_tool(name="predict_band_gap",
+#           description="Predict the band gap of a crystal structure using MEGNet multi-fidelity model from either a chemical formula or CIF file, with structure optimization")
+# async def predict_band_gap(
+#     formula: str = None,
+#     cif_file_name: str = None,
+#     method: str = "PBE",
+#     fmax: float = 0.01
+# ) -> str:
+#     """
+#     Predict the band gap of a crystal structure using the MEGNet multi-fidelity band gap model.
+
+#     First optimizes the crystal structure using M3GNet universal potential, then predicts
+#     the band gap based on the relaxed structure for more accurate results.
+
+#     Accepts either a chemical formula (searches Materials Project database) or a CIF file.
+
+#     Args:
+#         formula: Chemical formula to retrieve from Materials Project (e.g., "Fe2O3").
+#         cif_file_name: Name of CIF file in temp directory to use as structure source.
+#         method: The DFT method to use for the prediction. Options are "PBE", "GLLB-SC", "HSE", or "SCAN".
+#                 Default is "PBE".
+#         fmax: Maximum force tolerance for structure relaxation in eV/Å (default: 0.01).
+
+#     Returns:
+#         A string containing the predicted band gap in eV or an error message.
+#     """
+#     try:
+#         # First, relax the crystal structure
+#         relaxed_result = await relax_crystal_structure(
+#             formula=formula,
+#             cif_file_name=cif_file_name,
+#             fmax=fmax
+#         )
+
+#         # Check if relaxation was successful
+#         if isinstance(relaxed_result, str) and relaxed_result.startswith("Error"):
+#             return relaxed_result
+
+#         # Use the relaxed structure for band gap prediction
+#         structure = relaxed_result
+
+#         if structure is None:
+#             return "Error: Failed to obtain a valid relaxed structure"
+
+#         # Load the pre-trained MEGNet band gap model
+#         model = matgl.load_model("MEGNet-MP-2019.4.1-BandGap-mfi")
+
+#         # Map method name to index
+#         method_map = {"PBE": 0, "GLLB-SC": 1, "HSE": 2, "SCAN": 3}
+#         if method not in method_map:
+#             return f"Error: Unsupported method: {method}. Choose from PBE, GLLB-SC, HSE, or SCAN."
+
+#         # Set the graph label based on the method
+#         graph_attrs = torch.tensor([method_map[method]])
+
+#         # Predict the band gap using the relaxed structure
+#         bandgap = model.predict_structure(structure=structure, state_attr=graph_attrs)
+#         reduced_formula = structure.reduced_formula
+
+#         # Return the band gap as a string
+#         return f"The predicted band gap for relaxed {reduced_formula} using {method} method is {float(bandgap):.3f} eV."
+#     except Exception as e:
+#         return f"Error: {str(e)}"
+
+