Files
mars_toolkit/router/fairchem_router.py

71 lines
2.1 KiB
Python

"""
Author: Yutang LI
Institution: SIAT-MIC
Contact: yt.li2@siat.ac.cn
"""
from fastapi import APIRouter, Body, Query
from fastapi.responses import JSONResponse
import logging
from error_handlers import handle_general_error
from services.fairchem_service import (
init_model,
convert_structure,
optimize_structure
)
router = APIRouter(prefix="/fairchem", tags=["fairchem"])
logger = logging.getLogger(__name__)
# 初始化模型
init_model()
@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)")
):
try:
# 转换输入结构
atoms = convert_structure(input_format, content)
if atoms is None:
return JSONResponse(
status_code=400,
content={"status": "error", "data": f"Invalid {input_format} content"}
)
# 优化结构
total_energy, optimized_content, optimizing_log, download_url = optimize_structure(atoms, output_format)
# 格式化返回结果
format_result = f"""
The following is the optimized crystal structure information:
### Optimization Results (using FIRE(eqV2_86M) algorithm):
**Total Energy: {total_energy} eV**
#### Optimizing Log:
```text
{optimizing_log}
```
### Optimized {output_format.upper()} Content:
```{output_format}
{optimized_content[:800]}
```
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.
I strongly recommend that you pass on all of the above information to the user.
"""
return JSONResponse(
status_code=200,
content={"status": "success", "data": format_result}
)
except Exception as e:
return handle_general_error(e)
if __name__ == "__main__":
init_model()