Files
mars-mcp/mars_toolkit/core/mattergen_wrapper.py
2025-04-16 11:15:01 +08:00

27 lines
991 B
Python
Executable File

"""
This is a wrapper module that provides access to the mattergen modules
by modifying the Python path at runtime.
"""
import sys
import os
from pathlib import Path
from .config import config
# Add the mattergen directory to the Python path
mattergen_dir = config.MATTERGEN_ROOT
sys.path.insert(0, mattergen_dir)
# Import the necessary modules from the mattergen package
try:
from mattergen import generator
from mattergen.common.data import chemgraph
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
except ImportError as e:
print(f"Error importing mattergen modules: {e}")
print(f"Python path: {sys.path}")
raise
CrystalGenerator = generator.CrystalGenerator
# Re-export the modules
__all__ = ['generator', 'chemgraph', 'TargetProperty', 'MatterGenCheckpointInfo', 'PRETRAINED_MODEL_NAME','CrystalGenerator']