mattergen转服务

mattergen_client_example.py

@@ -0,0 +1,134 @@
+import requests
+import json
+import argparse
+import sys
+
+def generate_material(
+    url="http://localhost:8051/generate_material",
+    properties=None,
+    batch_size=2,
+    num_batches=1,
+    diffusion_guidance_factor=2.0
+):
+    """
+    调用MatterGen API生成晶体结构
+
+    Args:
+        url: API端点URL
+        properties: 可选的属性约束，例如{"dft_band_gap": 2.0}
+        batch_size: 每批生成的结构数量
+        num_batches: 批次数量
+        diffusion_guidance_factor: 控制生成结构与目标属性的符合程度
+
+    Returns:
+        生成的结构内容或错误信息
+    """
+    # 构建请求负载
+    payload = {
+        "properties": properties ,
+        "batch_size": batch_size,
+        "num_batches": num_batches,
+        "diffusion_guidance_factor": diffusion_guidance_factor
+    }
+
+    print(f"发送请求到 {url}")
+    print(f"请求参数: {json.dumps(payload, ensure_ascii=False, indent=2)}")
+
+    try:
+        # 添加headers参数，包含accept头
+        headers = {
+            "Content-Type": "application/json",
+            "accept": "application/json"
+        }
+
+        # 打印完整请求信息（调试用）
+        print(f"完整请求URL: {url}")
+        print(f"请求头: {headers}")
+        print(f"请求体: {json.dumps(payload)}")
+
+        # 禁用代理设置
+        proxies = {
+            "http": None,
+            "https": None
+        }
+
+        # 发送POST请求，添加headers参数，禁用代理，增加超时时间
+        response = requests.post(url, json=payload, headers=headers, proxies=proxies, timeout=300)
+
+        # 打印响应信息（调试用）
+        print(f"响应状态码: {response.status_code}")
+        print(f"响应头: {dict(response.headers)}")
+        print(f"响应内容: {response.text[:500]}...")  # 只打印前500个字符，避免输出过长
+
+        # 检查响应状态
+        if response.status_code == 200:
+            result = response.json()
+
+            if result["success"]:
+                print("\n生成成功!")
+                return result["content"]
+            else:
+                print(f"\n生成失败: {result['message']}")
+                return None
+        else:
+            print(f"\n请求失败，状态码: {response.status_code}")
+            print(f"响应内容: {response.text}")
+            return None
+
+    except Exception as e:
+        print(f"\n发生错误: {str(e)}")
+        print(f"错误类型: {type(e).__name__}")
+        import traceback
+        print(f"错误堆栈: {traceback.format_exc()}")
+        return None
+
+def main():
+    """命令行入口函数"""
+    parser = argparse.ArgumentParser(description="MatterGen API客户端示例")
+
+    # 添加命令行参数
+    parser.add_argument("--url", default="http://localhost:8051/generate_material",
+                        help="MatterGen API端点URL")
+    parser.add_argument("--property-name", default='dft_mag_density',help="属性名称，例如dft_band_gap")
+    parser.add_argument("--property-value",default=0.15,help="属性值，例如2.0")
+    parser.add_argument("--batch-size", type=int, default=2, help="每批生成的结构数量")
+    parser.add_argument("--num-batches", type=int, default=1, help="批次数量")
+    parser.add_argument("--guidance-factor", type=float, default=2.0,
+                        help="控制生成结构与目标属性的符合程度")
+
+    args = parser.parse_args()
+
+    # 构建属性字典
+    properties = None
+    if args.property_name and args.property_value:
+        try:
+            # 尝试将属性值转换为数字
+            try:
+                value = float(args.property_value)
+                # 如果是整数，转换为整数
+                if value.is_integer():
+                    value = int(value)
+            except ValueError:
+                # 如果无法转换为数字，保持为字符串
+                value = args.property_value
+
+            properties = {args.property_name: value}
+        except Exception as e:
+            print(f"解析属性值时出错: {str(e)}")
+            return
+
+    # 调用API
+    result = generate_material(
+        url=args.url,
+        properties=properties,
+        batch_size=args.batch_size,
+        num_batches=args.num_batches,
+        diffusion_guidance_factor=args.guidance_factor
+    )
+
+    if result:
+        print("\n生成的结构:")
+        print(result)
+
+if __name__ == "__main__":
+    main()