fairchem

constant.py

@@ -1,6 +1,13 @@
 TOPK_RESULT = 5
 TIME_OUT = 30
+
+# MP Configuration
 MP_API_KEY = "gfBp2in8qxm9Xm2SwLKFwNxDyZvNTAEt"
+MP_ENDPOINT = "https://api.materialsproject.org/"
+
+# Model
+FAIRCHEM_MODEL_PATH = "/home/ubuntu/sas0/LYT/mars1215/mars_toolkit/model/eqV2_86M_omat_mp_salex.pt"
+FMAX = 0.05
 
 # MinIO configuration
 MINIO_ENDPOINT = "https://s3-api.siat-mic.com"

database/material_project_router.py

@@ -11,7 +11,7 @@ import logging
 from mp_api.client import MPRester
 from multiprocessing import Process, Manager
 from typing import Dict, Any, List
-from constant import MP_API_KEY, TIME_OUT, TOPK_RESULT
+from constant import MP_ENDPOINT, MP_API_KEY, TIME_OUT, TOPK_RESULT
 
 router = APIRouter(prefix="/mp", tags=["Material Project"])
 
@@ -156,7 +156,7 @@ async def execute_search(search_args: Dict[str, Any], timeout: int = TIME_OUT) -
 def _search_worker(queue, api_key, **kwargs):
     """搜索工作线程"""
     try:
-        mpr = MPRester(api_key)
+        mpr = MPRester(api_key, endpoint=MP_ENDPOINT)
         result = mpr.materials.summary.search(**kwargs)
         queue.put([doc.model_dump() for doc in result])
     except Exception as e:

main.py

@@ -8,6 +8,7 @@ from fastapi import FastAPI
 import logging
 from database.material_project_router import router as material_project_router
 from database.oqmd_router import router as oqmd_router
+from model.fairchem_router import router as fairchem_router, init_model
 
 logging.basicConfig(
     level=logging.INFO,
@@ -16,5 +17,11 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
+
+@app.on_event("startup")
+def startup_event():
+    init_model()
+
 app.include_router(material_project_router)
 app.include_router(oqmd_router)
+app.include_router(fairchem_router)

model/fairchem_router.py

@@ -1,74 +1,165 @@
+from fastapi import APIRouter, Body, Query
 from fairchem.core import OCPCalculator
-from ase.optimize import FIRE # Import your optimizer of choice
-from ase.filters import FrechetCellFilter  # to include cell relaxations
-from ase.io import read
-from pymatgen.core import Structure
-from pymatgen.ext.matproj import MPRester
-from pymatgen.analysis.phase_diagram import PhaseDiagram, PDEntry
-from pymatgen.entries.computed_entries import ComputedStructureEntry
-from pymatgen.entries.mixing_scheme import MaterialsProjectDFTMixingScheme
+from ase.optimize import FIRE
+from ase.filters import FrechetCellFilter
+from ase.atoms import Atoms
+from ase.io import read, write
+from pymatgen.core.structure import Structure
+from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
+from pymatgen.io.cif import CifWriter
+import tempfile
+import os
+import boto3
+from constant import FAIRCHEM_MODEL_PATH, MINIO_ENDPOINT, MINIO_ACCESS_KEY, MINIO_SECRET_KEY, MINIO_BUCKET, INTERNEL_MINIO_ENDPOINT, FMAX
+from typing import Optional
+import logging
+import datetime
 
-# 创建相图并计算形成能与 above hull energy 的函数
-def calculate_phase_diagram_properties(structure, total_energy, api_key):
-    """
-    计算化合物的形成能和 above hull energy
-    参数:
-    - formula (str): 化学式 (如 "CsPbBr3")
-    - total_energy (float): 化合物的总能量 (eV)
-    - mpr (MPRester): MPRester 实例
+router = APIRouter(prefix="/fairchem", tags=["fairchem"])
+logger = logging.getLogger(__name__)
 
-    返回:
-    - formation_energy (float): 每个原子的形成能 (eV/atom)
-    - e_above_hull (float): 每个原子的 above hull energy (eV/atom)
-    """
-    chemsys = structure.chemical_system.split("-")
-    formula = structure.reduced_formula
-    with MPRester(api_key) as mpr:
-        # 获取化学系统中所有的相
-        # entries = mpr.get_entries_in_chemsys(elements=chemsys, additional_criteria={"thermo_types": ["GGA_GGA+U"]})
-        entries = mpr.get_entries_in_chemsys(elements=chemsys, additional_criteria={"thermo_types": ["GGA_GGA+U", "R2SCAN"]})
+# 初始化模型
+calc = None
 
-    # 创建新计算结构的 PDEntry
-    pd_entry = PDEntry(composition=formula, energy=total_energy)
-    # entries.append(pd_entry)
+def init_model():
+    global calc
+    calc = OCPCalculator(checkpoint_path=FAIRCHEM_MODEL_PATH)
 
-    scheme = MaterialsProjectDFTMixingScheme()
-    entries = scheme.process_entries(entries)
+def convert_structure(input_format: str, content: str) -> Optional[Atoms]:
+    """将输入内容转换为Atoms对象"""
+    try:
+        with tempfile.NamedTemporaryFile(suffix=f".{input_format}", mode="w", delete=False) as tmp_file:
+            tmp_file.write(content)
+            tmp_path = tmp_file.name
         
-    # 创建相图
-    pd = PhaseDiagram(entries + [pd_entry])
+        atoms = read(tmp_path)
+        os.unlink(tmp_path)
+        return atoms
+    except Exception as e:
+        logger.error(f"Failed to convert structure: {str(e)}")
+        return None
 
-    # 计算形成能和 above hull energy
-    formation_energy = pd.get_form_energy_per_atom(pd_entry)
-    e_above_hull = pd.get_e_above_hull(pd_entry)
+def generate_symmetry_cif(structure: Structure) -> str:
+    """生成对称性CIF"""
+    analyzer = SpacegroupAnalyzer(structure)
+    structure = analyzer.get_refined_structure()
     
-    return formation_energy, e_above_hull
+    with tempfile.NamedTemporaryFile(suffix=".cif", mode="w+", delete=False) as tmp_file:
+        cif_writer = CifWriter(structure, symprec=0.1, refine_struct=True)
+        cif_writer.write_file(tmp_file.name)
+        tmp_file.seek(0)
+        return tmp_file.read()
 
+def upload_to_minio(file_path: str, file_name: str) -> str:
+    """上传文件到MinIO并返回预签名URL"""
+    try:
+        minio_client = boto3.client(
+            's3',
+            endpoint_url=MINIO_ENDPOINT if INTERNEL_MINIO_ENDPOINT == "" else INTERNEL_MINIO_ENDPOINT,
+            aws_access_key_id=MINIO_ACCESS_KEY,
+            aws_secret_access_key=MINIO_SECRET_KEY
+        )
         
-atoms = read("/home/ubuntu/sas0/LYT/paper_dataset/mp_cif/CsPbBr3.cif") # Read in an atoms object or create your own structure
-calc = OCPCalculator(checkpoint_path="/home/ubuntu/sas0/LYT/paper_dataset/mp_cif/meta_fairchem/eqV2_86M_omat_mp_salex.pt") # Path to downloaded checkpoint
-atoms.calc = calc
-dyn = FIRE(FrechetCellFilter(atoms))
-dyn.run(fmax=0.01)
+        bucket_name = MINIO_BUCKET
+        minio_client.upload_file(file_path, bucket_name, file_name, ExtraArgs={"ACL": "private"})
         
-total_energy = atoms.get_potential_energy()
-print("Predicted Total Energy: ", total_energy)
+        # 生成预签名 URL
+        url = minio_client.generate_presigned_url(
+            'get_object',
+            Params={'Bucket': bucket_name, 'Key': file_name},
+            ExpiresIn=3600
+        )
+        return url.replace(INTERNEL_MINIO_ENDPOINT, MINIO_ENDPOINT)
+    except Exception as e:
+        logger.error(f"Failed to upload to MinIO: {str(e)}")
+        raise RuntimeError(f"MinIO upload failed: {str(e)}") from e
 
-# 保存优化后的结构
-atoms.write("optimized_structure.cif")  # 保存为 CIF 文件
-print("Geometry optimization completed. Optimized structure saved as 'optimized_structure.cif'.")
+from io import StringIO
+import sys
 
-# 从 ASE 转换为 Pymatgen 结构
-optimized_structure = Structure.from_file("optimized_structure.cif")
+def optimize_structure(atoms: Atoms, output_format: str):
+    """优化晶体结构"""
+    atoms.calc = calc
     
-api_key = "gfBp2in8qxm9Xm2SwLKFwNxDyZvNTAEt"
-mpr = MPRester(api_key)
-print(f"Chemical Formula: {optimized_structure .composition.reduced_formula}")
-formation_energy, e_above_hull = calculate_phase_diagram_properties(
-    structure=optimized_structure,
-    total_energy=total_energy,
-    api_key=api_key
-)
+    # 捕获优化日志
+    old_stdout = sys.stdout
+    sys.stdout = log_capture = StringIO()
     
-print(formation_energy, e_above_hull)
-print()
+    try:
+        dyn = FIRE(FrechetCellFilter(atoms))
+        dyn.run(fmax=FMAX)
+        total_energy = atoms.get_total_energy()
+        optimization_log = log_capture.getvalue()
+    finally:
+        sys.stdout = old_stdout
+
+    # 处理对称性
+    if output_format == "cif":
+        optimized_structure = Structure.from_ase_atoms(atoms)
+        content = generate_symmetry_cif(optimized_structure)
+    else:
+        with tempfile.NamedTemporaryFile(suffix=f".{output_format}", mode="w+", delete=False) as tmp_file:
+            write(tmp_file.name, atoms)
+            tmp_file.seek(0)
+            content = tmp_file.read()
+    
+    # 保存优化结果到临时文件
+    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+    file_name = f"optimized_structure_{timestamp}.{output_format}"
+    with tempfile.NamedTemporaryFile(suffix=f".{output_format}", mode="w", delete=False) as tmp_file:
+        tmp_file.write(content)
+        tmp_path = tmp_file.name
+    try:
+        # 上传到MinIO
+        url = upload_to_minio(tmp_path, file_name)
+        return total_energy, content, url, optimization_log
+    finally:
+        os.unlink(tmp_path)
+
+@router.post("/optimize_structure")
+async def optimize_structure_endpoint(
+    content: str = Body(..., description="Input structure content"),
+    input_format: str = Query("cif", description="Input format (cif, poscar, json, xyz)"),
+    output_format: str = Query("cif", description="Output format (cif, poscar, json, xyz)")
+):
+    # 转换输入结构
+    atoms = convert_structure(input_format, content)
+    if atoms is None:
+        return {
+            "status": "error",
+            "data": f"Invalid {input_format} content"
+        }
+    
+    try:
+        # 优化结构
+        total_energy, optimized_content, download_url, optimization_log = optimize_structure(atoms, output_format)
+        
+        # 格式化返回结果
+        format_result = f"""
+The following is the optimized crystal structure information:
+
+### Optimization Results (using FIRE(eqV2_86M) algorithm):
+```text
+{optimization_log}
+```
+Finally, the Total Energy is: {total_energy} eV
+Due to length limitations, the complete {output_format.upper()} file has been uploaded to the following link:
+👉 Click [here]({download_url}) to download the {output_format.upper()} file
+
+Please ensure that the Optimization Results and download link are fully conveyed to the user, as this is very important for them.
+"""
+        print(format_result)
+        return {
+            "status": "success",
+            "data": format_result
+        }
+    
+    except Exception as e:
+        logger.error(f"Optimization failed: {str(e)}")
+        return {
+            "status": "error",
+            "data": str(e)
+        }
+
+if __name__ == "__main__":
+    init_model()