mattergen转服务

mars_toolkit/services/__init__.py

@@ -0,0 +1,12 @@
+"""
+Services module for mars_toolkit.
+
+This module contains service classes that provide persistent functionality
+across multiple function calls, such as maintaining initialized models.
+"""
+
+# Import services for easy access
+from mars_toolkit.services.mattergen_service import MatterGenService
+
+# Export services
+__all__ = ['MatterGenService']

mars_toolkit/services/mattergen_service.py

@@ -0,0 +1,342 @@
+"""
+MatterGen service for mars_toolkit.
+
+This module provides a service for generating crystal structures using MatterGen.
+The service initializes the CrystalGenerator once and reuses it for multiple
+generation requests, improving performance.
+"""
+
+import os
+import logging
+import json
+from pathlib import Path
+from typing import Dict, Any, Optional, Union, List
+import threading
+
+# 导入mattergen相关模块
+import sys
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
+from mattergen_wrapper import generator
+CrystalGenerator = generator.CrystalGenerator
+from mattergen.common.data.types import TargetProperty
+from mattergen.common.utils.eval_utils import MatterGenCheckpointInfo
+from mattergen.common.utils.data_classes import PRETRAINED_MODEL_NAME
+
+# 导入mars_toolkit配置
+from mars_toolkit.core.config import config
+
+logger = logging.getLogger(__name__)
+
+class MatterGenService:
+    """
+    Service for generating crystal structures using MatterGen.
+
+    This service initializes the CrystalGenerator once and reuses it for multiple
+    generation requests, improving performance.
+    """
+
+    _instance = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls):
+        """
+        Get the singleton instance of MatterGenService.
+
+        Returns:
+            MatterGenService: The singleton instance.
+        """
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        """
+        Initialize the MatterGenService.
+
+        This initializes the base generator without any property conditioning.
+        Specific generators for different property conditions will be initialized
+        on demand.
+        """
+        self._generators = {}
+        self._output_dir = config.MATTERGENMODEL_RESULT_PATH
+
+        # 确保输出目录存在
+        if not os.path.exists(self._output_dir):
+            os.makedirs(self._output_dir)
+
+        # 初始化基础生成器（无条件生成）
+        self._init_base_generator()
+
+    def _init_base_generator(self):
+        """
+        Initialize the base generator for unconditional generation.
+        """
+        model_path = os.path.join(config.MATTERGENMODEL_ROOT, "mattergen_base")
+
+        if not os.path.exists(model_path):
+            logger.warning(f"Base model directory not found at {model_path}. MatterGen service may not work properly.")
+            return
+
+        logger.info(f"Initializing base MatterGen generator from {model_path}")
+
+        try:
+            checkpoint_info = MatterGenCheckpointInfo(
+                model_path=Path(model_path).resolve(),
+                load_epoch="last",
+                config_overrides=[],
+                strict_checkpoint_loading=True,
+            )
+
+            generator = CrystalGenerator(
+                checkpoint_info=checkpoint_info,
+                properties_to_condition_on=None,
+                batch_size=2,  # 默认值，可在生成时覆盖
+                num_batches=1,  # 默认值，可在生成时覆盖
+                sampling_config_name="default",
+                sampling_config_path=None,
+                sampling_config_overrides=[],
+                record_trajectories=True,
+                diffusion_guidance_factor=0.0,
+                target_compositions_dict=[],
+            )
+
+            self._generators["base"] = generator
+            logger.info("Base MatterGen generator initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize base MatterGen generator: {e}")
+
+    def _get_or_create_generator(
+        self,
+        properties: Optional[Dict[str, Any]] = None,
+        batch_size: int = 2,
+        num_batches: int = 1,
+        diffusion_guidance_factor: float = 2.0
+    ):
+        """
+        Get or create a generator for the specified properties.
+
+        Args:
+            properties: Optional property constraints
+            batch_size: Number of structures per batch
+            num_batches: Number of batches to generate
+            diffusion_guidance_factor: Controls adherence to target properties
+
+        Returns:
+            tuple: (generator, generator_key, properties_to_condition_on)
+        """
+        # 如果没有属性约束，使用基础生成器
+        if not properties:
+            if "base" not in self._generators:
+                self._init_base_generator()
+            return self._generators.get("base"), "base", None
+
+        # 处理属性约束
+        properties_to_condition_on = {}
+        for property_name, property_value in properties.items():
+            properties_to_condition_on[property_name] = property_value
+
+        # 确定模型目录
+        if len(properties) == 1:
+            # 单属性条件
+            property_name = list(properties.keys())[0]
+            property_to_model = {
+                "dft_mag_density": "dft_mag_density",
+                "dft_bulk_modulus": "dft_bulk_modulus",
+                "dft_shear_modulus": "dft_shear_modulus",
+                "energy_above_hull": "energy_above_hull",
+                "formation_energy_per_atom": "formation_energy_per_atom",
+                "space_group": "space_group",
+                "hhi_score": "hhi_score",
+                "ml_bulk_modulus": "ml_bulk_modulus",
+                "chemical_system": "chemical_system",
+                "dft_band_gap": "dft_band_gap"
+            }
+            model_dir = property_to_model.get(property_name, property_name)
+            generator_key = f"single_{property_name}"
+        else:
+            # 多属性条件
+            property_keys = set(properties.keys())
+            if property_keys == {"dft_mag_density", "hhi_score"}:
+                model_dir = "dft_mag_density_hhi_score"
+                generator_key = "multi_dft_mag_density_hhi_score"
+            elif property_keys == {"chemical_system", "energy_above_hull"}:
+                model_dir = "chemical_system_energy_above_hull"
+                generator_key = "multi_chemical_system_energy_above_hull"
+            else:
+                # 如果没有特定的多属性模型，使用第一个属性的模型
+                first_property = list(properties.keys())[0]
+                model_dir = first_property
+                generator_key = f"multi_{first_property}_etc"
+
+        # 构建完整的模型路径
+        model_path = os.path.join(config.MATTERGENMODEL_ROOT, model_dir)
+
+        # 检查模型目录是否存在
+        if not os.path.exists(model_path):
+            # 如果特定模型不存在，回退到基础模型
+            logger.warning(f"Model directory for {model_dir} not found. Using base model instead.")
+            model_path = os.path.join(config.MATTERGENMODEL_ROOT, "mattergen_base")
+            generator_key = "base"
+
+        # 检查是否已经有这个生成器
+        if generator_key in self._generators:
+            # 更新生成器的参数
+            generator = self._generators[generator_key]
+            generator.batch_size = batch_size
+            generator.num_batches = num_batches
+            generator.diffusion_guidance_factor = diffusion_guidance_factor if properties else 0.0
+            return generator, generator_key, properties_to_condition_on
+
+        # 创建新的生成器
+        try:
+            logger.info(f"Initializing new MatterGen generator for {generator_key} from {model_path}")
+
+            checkpoint_info = MatterGenCheckpointInfo(
+                model_path=Path(model_path).resolve(),
+                load_epoch="last",
+                config_overrides=[],
+                strict_checkpoint_loading=True,
+            )
+
+            generator = CrystalGenerator(
+                checkpoint_info=checkpoint_info,
+                properties_to_condition_on=properties_to_condition_on,
+                batch_size=batch_size,
+                num_batches=num_batches,
+                sampling_config_name="default",
+                sampling_config_path=None,
+                sampling_config_overrides=[],
+                record_trajectories=True,
+                diffusion_guidance_factor=diffusion_guidance_factor if properties else 0.0,
+                target_compositions_dict=[],
+            )
+
+            self._generators[generator_key] = generator
+            logger.info(f"MatterGen generator for {generator_key} initialized successfully")
+            return generator, generator_key, properties_to_condition_on
+        except Exception as e:
+            logger.error(f"Failed to initialize MatterGen generator for {generator_key}: {e}")
+            # 回退到基础生成器
+            if "base" not in self._generators:
+                self._init_base_generator()
+            return self._generators.get("base"), "base", None
+
+    def generate(
+        self,
+        properties: Optional[Dict[str, Union[float, str, Dict[str, Union[float, str]]]]] = None,
+        batch_size: int = 2,
+        num_batches: int = 1,
+        diffusion_guidance_factor: float = 2.0
+    ) -> str:
+        """
+        Generate crystal structures with optional property constraints.
+
+        Args:
+            properties: Optional property constraints
+            batch_size: Number of structures per batch
+            num_batches: Number of batches to generate
+            diffusion_guidance_factor: Controls adherence to target properties
+
+        Returns:
+            str: Descriptive text with generated crystal structures in CIF format
+        """
+        from mars_toolkit.compute.material_gen import format_cif_content
+
+        # 处理字符串输入（如果提供）
+        if isinstance(properties, str):
+            try:
+                properties = json.loads(properties)
+            except json.JSONDecodeError:
+                raise ValueError(f"Invalid properties JSON string: {properties}")
+
+        # 如果为None，默认为空字典
+        properties = properties or {}
+
+        # 获取或创建生成器
+        generator, generator_key, properties_to_condition_on = self._get_or_create_generator(
+            properties, batch_size, num_batches, diffusion_guidance_factor
+        )
+
+        if generator is None:
+            return "Error: Failed to initialize MatterGen generator"
+
+        # 生成结构
+        try:
+            generator.generate(output_dir=Path(self._output_dir))
+        except Exception as e:
+            logger.error(f"Error generating structures: {e}")
+            return f"Error generating structures: {e}"
+
+        # 创建字典存储文件内容
+        result_dict = {}
+
+        # 定义文件路径
+        cif_zip_path = os.path.join(self._output_dir, "generated_crystals_cif.zip")
+        xyz_file_path = os.path.join(self._output_dir, "generated_crystals.extxyz")
+        trajectories_zip_path = os.path.join(self._output_dir, "generated_trajectories.zip")
+
+        # 读取CIF压缩文件
+        if os.path.exists(cif_zip_path):
+            with open(cif_zip_path, 'rb') as f:
+                result_dict['cif_content'] = f.read()
+
+        # 根据生成类型创建描述性提示
+        if not properties:
+            generation_type = "unconditional"
+            title = "Generated Material Structures"
+            description = "These structures were generated unconditionally, meaning no specific properties were targeted."
+            property_description = "unconditionally"
+        elif len(properties) == 1:
+            generation_type = "single_property"
+            property_name = list(properties.keys())[0]
+            property_value = properties[property_name]
+            title = f"Generated Material Structures Conditioned on {property_name} = {property_value}"
+            description = f"These structures were generated with property conditioning, targeting a {property_name} value of {property_value}."
+            property_description = f"conditioned on {property_name} = {property_value}"
+        else:
+            generation_type = "multi_property"
+            title = "Generated Material Structures Conditioned on Multiple Properties"
+            description = "These structures were generated with multi-property conditioning, targeting the specified property values."
+            property_description = f"conditioned on multiple properties: {', '.join([f'{name} = {value}' for name, value in properties.items()])}"
+
+        # 创建完整的提示
+        prompt = f"""
+# {title}
+
+This data contains {batch_size * num_batches} crystal structures generated by the MatterGen model, {property_description}.
+
+    {'' if generation_type == 'unconditional' else f'''
+    A diffusion guidance factor of {diffusion_guidance_factor} was used, which controls how strongly
+    the generation adheres to the specified property values. Higher values produce samples that more
+    closely match the target properties but may reduce diversity.
+    '''}
+
+## CIF Files (Crystallographic Information Files)
+
+- Standard format for crystallographic structures
+- Contains unit cell parameters, atomic positions, and symmetry information
+- Used by crystallographic software and visualization tools
+
+```
+{format_cif_content(result_dict.get('cif_content', b'').decode('utf-8', errors='replace') if isinstance(result_dict.get('cif_content', b''), bytes) else str(result_dict.get('cif_content', '')))}
+```
+
+{description}
+You can use these structures for materials discovery, property prediction, or further analysis.
+"""
+
+        # 清理文件（读取后删除）
+        try:
+            if os.path.exists(cif_zip_path):
+                os.remove(cif_zip_path)
+            if os.path.exists(xyz_file_path):
+                os.remove(xyz_file_path)
+            if os.path.exists(trajectories_zip_path):
+                os.remove(trajectories_zip_path)
+        except Exception as e:
+            logger.warning(f"Error cleaning up files: {e}")
+
+        return prompt