488 lines
22 KiB
Python
488 lines
22 KiB
Python
"""
|
||
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)}"
|
||
|
||
|