multi_mcp/sci_mcp/material_mcp/mattergen_gen/mattergen_service.py

"""
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 datetime
import os
import logging
import json
from pathlib import Path
import re
from typing import Dict, Any, Optional, Union, List
import threading
import torch

from .mattergen_wrapper import *
from ...core.config import material_config

logger = logging.getLogger(__name__)



def format_cif_content(content):
    """
    Format CIF content by removing unnecessary headers and organizing each CIF file.

    Args:
        content: String containing CIF content, possibly with PK headers

    Returns:
        Formatted string with each CIF file properly labeled and formatted
    """
    # 如果内容为空，直接返回空字符串
    if not content or content.strip() == '':
        return ''

    # 删除从PK开始到第一个_chemical_formula_structural之前的所有内容
    content = re.sub(r'PK.*?(?=_chemical_formula_structural)', '', content, flags=re.DOTALL)

    # 删除从PK开始到字符串结束且没有_chemical_formula_structural的内容
    content = re.sub(r'PK[^_]*$', '', content, flags=re.DOTALL)
    content = re.sub(r'PK.*?(?!.*_chemical_formula_structural)$', '', content, flags=re.DOTALL)

    # 使用_chemical_formula_structural作为分隔符来分割不同的CIF文件
    # 但我们需要保留这个字段在每个CIF文件中
    cif_blocks = []

    # 查找所有_chemical_formula_structural的位置
    formula_positions = [m.start() for m in re.finditer(r'_chemical_formula_structural', content)]

    # 如果没有找到任何_chemical_formula_structural，返回空字符串
    if not formula_positions:
        return ''

    # 分割CIF块
    for i in range(len(formula_positions)):
        start_pos = formula_positions[i]
        # 如果是最后一个块，结束位置是字符串末尾
        end_pos = formula_positions[i+1] if i < len(formula_positions)-1 else len(content)

        cif_block = content[start_pos:end_pos].strip()

        # 提取formula值
        formula_match = re.search(r'_chemical_formula_structural\s+(\S+)', cif_block)
        if formula_match:
            formula = formula_match.group(1)
            cif_blocks.append((formula, cif_block))

    # 格式化输出
    result = []
    for i, (formula, cif_content) in enumerate(cif_blocks, 1):
        formatted = f"[cif {i} begin]\ndata_{formula}\n{cif_content}\n[cif {i} end]\n"
        result.append(formatted)

    return "\n".join(result)

def extract_cif_file_from_zip(cifs_zip_path: str):
    """
    Extract CIF files from a zip archive, extract formula from each CIF file,
    and save each CIF file with its formula as the filename.

    Args:
        cifs_zip_path: Path to the zip file

    Returns:
        list: List of tuples containing (index, formula, cif_path)
    """
    result_dict = {}
    if os.path.exists(cifs_zip_path):
            with open(cifs_zip_path, 'rb') as f:
                result_dict['cif_content'] = f.read()
    cifs_content = 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', '')))
    pattern = r'\[cif (\d+) begin\]\n(.*?)\n\[cif \1 end\]'
    matches = re.findall(pattern, cifs_content, re.DOTALL)

    # 处理每个匹配项，提取formula并保存CIF文件
    saved_files = []
    for idx, cif_content in matches:
        # 提取data_{formula}中的formula
        formula_match = re.search(r'data_([^\s]+)', cif_content)
        if formula_match:
            formula = formula_match.group(1)
            # 构建保存路径
            cif_path = os.path.join(material_config.TEMP_ROOT, f"{formula}.cif")
            # 保存CIF文件
            with open(cif_path, 'w') as f:
                f.write(cif_content)
            saved_files.append((idx, formula, cif_path))

    return saved_files


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()

    # 模型到GPU ID的映射
    MODEL_TO_GPU = {
        "mattergen_base": "0",                          # 基础模型使用GPU 0
        "dft_mag_density": "1",                         # 磁密度模型使用GPU 1
        "dft_bulk_modulus": "2",                        # 体积模量模型使用GPU 2
        "dft_shear_modulus": "3",                       # 剪切模量模型使用GPU 3
        "energy_above_hull": "4",                       # 能量模型使用GPU 4
        "formation_energy_per_atom": "5",               # 形成能模型使用GPU 5
        "space_group": "6",                             # 空间群模型使用GPU 6
        "hhi_score": "7",                               # HHI评分模型使用GPU 7
        "ml_bulk_modulus": "0",                         # ML体积模量模型使用GPU 0
        "chemical_system": "1",                         # 化学系统模型使用GPU 1
        "dft_band_gap": "2",                            # 带隙模型使用GPU 2
        "dft_mag_density_hhi_score": "3",               # 多属性模型使用GPU 3
        "chemical_system_energy_above_hull": "4"        # 多属性模型使用GPU 4
    }

    @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 = material_config.TEMP_ROOT

        # 确保输出目录存在
        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(material_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, gpu_id)
        """
        # 如果没有属性约束，使用基础生成器
        if not properties:
            if "base" not in self._generators:
                self._init_base_generator()
            gpu_id = self.MODEL_TO_GPU.get("mattergen_base", "0")  # 默认使用GPU 0
            return self._generators.get("base"), "base", None, gpu_id

        # 处理属性约束
        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"

        # 获取对应的GPU ID
        gpu_id = self.MODEL_TO_GPU.get(model_dir, "0")  # 默认使用GPU 0

        # 构建完整的模型路径
        model_path = os.path.join(material_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(material_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, gpu_id

        # 创建新的生成器
        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, gpu_id
        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()
            base_gpu_id = self.MODEL_TO_GPU.get("mattergen_base", "0")
            return self._generators.get("base"), "base", None, base_gpu_id

    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
        """


        # 处理字符串输入（如果提供）
        if isinstance(properties, str):
            try:
                properties = json.loads(properties)
            except json.JSONDecodeError:
                raise ValueError(f"Invalid properties JSON string: {properties}")

        # 如果为None，默认为空字典
        properties = properties or {}

        # 获取或创建生成器和GPU ID
        generator, generator_key, properties_to_condition_on, gpu_id = self._get_or_create_generator(
            properties, batch_size, num_batches, diffusion_guidance_factor
        )
        print("gpu_id",gpu_id)
        if generator is None:
            return "Error: Failed to initialize MatterGen generator"

        # 使用torch.cuda.set_device()直接设置当前GPU
        try:
            # 将字符串类型的gpu_id转换为整数
            cuda_device_id = int(gpu_id)
            torch.cuda.set_device(cuda_device_id)
            logger.info(f"Setting CUDA device to GPU {cuda_device_id} for model {generator_key}")
            print(f"Using GPU {cuda_device_id} (CUDA device index) for model {generator_key}")
        except Exception as e:
            logger.warning(f"Error setting CUDA device: {e}. Falling back to default device.")

        # 生成结构
        try:

            output_dir= Path(self._output_dir+f'/{datetime.datetime.now().strftime("%Y%m%d%H%M%S")}')
            Path.mkdir(output_dir, parents=True, exist_ok=True)
            generator.generate(output_dir=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(str(output_dir), f"generated_crystals_cif.zip")
        xyz_file_path = os.path.join(str(output_dir), f"generated_crystals.extxyz")
        trajectories_zip_path = os.path.join(str(output_dir), f"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.
"""
        # print("prompt",prompt)
        # 清理文件（读取后删除）
        # 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}")

        # GPU设备已经在生成前由torch.cuda.set_device()设置，不需要额外清理
        logger.info(f"Generation completed on GPU for model {generator_key}")

        return prompt