DynamisLab/src/config.py
2026-02-20 11:57:01 +08:00

126 lines
3.4 KiB
Python

"""
Configuration management for DynamisLab.
Handles loading of CelerisLab configurations and project settings.
"""
import os
import sys
from pathlib import Path
from typing import Optional, Tuple
# Determine project root directory
_current_file = Path(__file__).resolve()
_project_root = _current_file.parent.parent.parent # Go up to DynamisLabNew/
# Configuration directory (relative to project root)
CONFIG_DIR = _project_root / 'configs'
# Output directories
MODELS_DIR = _project_root / 'models'
OUTPUT_DIR = _project_root / 'output'
TENSORBOARD_DIR = _project_root / 'tensorboard'
# Create output directories if they don't exist
MODELS_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)
TENSORBOARD_DIR.mkdir(exist_ok=True)
def setup_celeris_import() -> None:
"""
Setup CelerisLab import path.
Assumes CelerisLab is either:
1. Installed as a package (pip install CelerisLab)
2. Available as a git submodule in project root
3. Available via PYTHONPATH environment variable
"""
try:
# Try to import CelerisLab (it might be installed)
import CelerisLab
return
except ImportError:
pass
# Check for CelerisLab as submodule
celeris_submodule = _project_root / 'CelerisLab' / 'src'
if celeris_submodule.exists():
sys.path.insert(0, str(celeris_submodule))
return
# If still not found, raise error with helpful message
raise ImportError(
"CelerisLab not found. Please either:\n"
" 1. Install it: pip install -e ../CelerisLab\n"
" 2. Add as git submodule: git submodule add <url> CelerisLab\n"
" 3. Set PYTHONPATH to include CelerisLab src directory"
)
def load_celeris_configs(
cuda_config_path: Optional[str] = None,
field_config_path: Optional[str] = None
) -> Tuple:
"""
Load CelerisLab configurations.
Args:
cuda_config_path: Optional path to CUDA config. If None, uses CONFIG_DIR.
field_config_path: Optional path to field config. If None, uses CONFIG_DIR.
Returns:
Tuple of (config_cuda, config_field)
"""
# Setup CelerisLab import
setup_celeris_import()
from CelerisLab import utils
# Set environment variable to point to our configs
os.environ['CELERISLAB_CONFIG_DIR'] = str(CONFIG_DIR)
# Load configurations - CelerisLab will find them automatically
if cuda_config_path is None:
config_cuda = utils.load_cuda_config()
else:
config_cuda = utils.load_cuda_config(cuda_config_path)
if field_config_path is None:
config_field = utils.load_flow_field_config()
else:
config_field = utils.load_flow_field_config(field_config_path)
return config_cuda, config_field
def get_model_path(model_name: str) -> Path:
"""Get full path for a model file."""
return MODELS_DIR / f"{model_name}.zip"
def get_tensorboard_logdir(run_name: str) -> Path:
"""Get TensorBoard log directory for a run."""
logdir = TENSORBOARD_DIR / run_name
logdir.mkdir(exist_ok=True)
return logdir
def get_output_path(filename: str) -> Path:
"""Get full path for an output file."""
return OUTPUT_DIR / filename
# Expose project paths for convenience
__all__ = [
'CONFIG_DIR',
'MODELS_DIR',
'OUTPUT_DIR',
'TENSORBOARD_DIR',
'setup_celeris_import',
'load_celeris_configs',
'get_model_path',
'get_tensorboard_logdir',
'get_output_path',
]