First commit
This commit is contained in:
Frank14f
@@ -0,0 +1,11 @@
|
||||
"""
|
||||
DynamisLab: Machine Learning for Computational Fluid Dynamics
|
||||
"""
|
||||
|
||||
__version__ = '0.1.0'
|
||||
__author__ = 'Frank14f'
|
||||
|
||||
from . import config
|
||||
from . import environments
|
||||
|
||||
__all__ = ['config', 'environments', '__version__']
|
||||
+125
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
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',
|
||||
]
|
||||
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
Gymnasium environments for CFD control tasks.
|
||||
"""
|
||||
|
||||
from .cfd_env import CFDFlowControlEnv
|
||||
|
||||
__all__ = ['CFDFlowControlEnv']
|
||||
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
CFD Flow Control Environment using CelerisLab.
|
||||
|
||||
A Gymnasium environment for active flow control using lattice Boltzmann simulation.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections import deque
|
||||
from typing import Optional, Tuple, Dict, Any
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
from gymnasium import spaces
|
||||
|
||||
# Set threading to avoid conflicts with GPU
|
||||
os.environ["OMP_NUM_THREADS"] = "1"
|
||||
os.environ["MKL_NUM_THREADS"] = "1"
|
||||
|
||||
|
||||
class CFDFlowControlEnv(gym.Env):
|
||||
"""
|
||||
CFD flow control environment with cylinder and sensors.
|
||||
|
||||
The environment simulates flow around a cylinder with multiple control cylinders
|
||||
and sensors to measure flow properties. The agent controls the cylinder velocities
|
||||
to optimize flow characteristics.
|
||||
|
||||
Args:
|
||||
device_id: CUDA device ID to use for simulation
|
||||
config_cuda: CelerisLab CUDA configuration (optional, will load from config if None)
|
||||
config_field: CelerisLab flow field configuration (optional, will load from config if None)
|
||||
n_control_cylinders: Number of controllable cylinders (default: 3)
|
||||
n_sensors: Number of flow sensors (default: 3)
|
||||
max_steps: Maximum steps per episode (default: 500)
|
||||
sample_interval: Simulation steps between observations (default: 800)
|
||||
fifo_length: Length of state history (default: 120)
|
||||
convergence_length: Steps to check for convergence (default: 60)
|
||||
warmup_steps_factor: Multiple of grid size for warmup (default: 4)
|
||||
"""
|
||||
|
||||
metadata = {"render_modes": ["human"], "render_fps": 30}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device_id: int = 0,
|
||||
config_cuda = None,
|
||||
config_field = None,
|
||||
n_control_cylinders: int = 3,
|
||||
n_sensors: int = 3,
|
||||
max_steps: int = 500,
|
||||
sample_interval: int = 800,
|
||||
fifo_length: int = 120,
|
||||
convergence_length: int = 60,
|
||||
warmup_steps_factor: int = 4,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Load configurations if not provided
|
||||
if config_cuda is None or config_field is None:
|
||||
from ..config import load_celeris_configs
|
||||
config_cuda, config_field = load_celeris_configs()
|
||||
|
||||
self.config_cuda = config_cuda
|
||||
self.config_field = config_field
|
||||
self.device_id = device_id
|
||||
|
||||
# Environment parameters
|
||||
self.n_control = n_control_cylinders
|
||||
self.n_sensors = n_sensors
|
||||
self.max_steps = max_steps
|
||||
self.sample_interval = sample_interval
|
||||
self.fifo_length = fifo_length
|
||||
self.convergence_length = convergence_length
|
||||
self.warmup_steps_factor = warmup_steps_factor
|
||||
|
||||
# Determine data type
|
||||
if config_field.data_type == "FP32":
|
||||
self.dtype = np.float32
|
||||
elif config_field.data_type == "FP64":
|
||||
self.dtype = np.float64
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {config_field.data_type}")
|
||||
|
||||
# Action and observation dimensions
|
||||
# Action: velocity control for n cylinders (x, y, rotation)
|
||||
self.action_dim = n_control_cylinders
|
||||
# Observation: sensor readings (u, v) from n sensors
|
||||
self.obs_dim = n_sensors * 2 * 2 # 2 velocity components × 2 (current + derivative)
|
||||
|
||||
# Gym spaces
|
||||
self.action_space = spaces.Box(
|
||||
low=-1.0,
|
||||
high=1.0,
|
||||
shape=(self.action_dim,),
|
||||
dtype=self.dtype
|
||||
)
|
||||
self.observation_space = spaces.Box(
|
||||
low=-np.inf,
|
||||
high=np.inf,
|
||||
shape=(self.obs_dim,),
|
||||
dtype=self.dtype
|
||||
)
|
||||
|
||||
# State tracking
|
||||
self.fifo_states = deque(maxlen=fifo_length)
|
||||
self.target_states = np.empty((0, self.n_sensors * 2), dtype=self.dtype)
|
||||
self.current_step = 0
|
||||
|
||||
# Normalization factors (will be set during warmup)
|
||||
self.sens_norm_fact = np.ones(self.n_sensors * 2, dtype=self.dtype)
|
||||
self.sens_deviation = np.zeros(self.n_sensors * 2, dtype=self.dtype)
|
||||
|
||||
# Reward tracking
|
||||
self.reward_cd = 0.0
|
||||
self.reward_cl = 0.0
|
||||
self.reward_sim = 0.0
|
||||
|
||||
# Initialize flow field
|
||||
self._init_flow_field()
|
||||
|
||||
def _init_flow_field(self):
|
||||
"""Initialize the CelerisLab flow field simulation."""
|
||||
from CelerisLab import FlowField
|
||||
|
||||
self.flow_field = FlowField(
|
||||
self.config_field,
|
||||
self.config_cuda,
|
||||
self.device_id
|
||||
)
|
||||
|
||||
# Get grid parameters
|
||||
L0 = 20 # Characteristic length
|
||||
U0 = self.config_field.velocity
|
||||
NX = self.flow_field.FIELD_SHAPE[0]
|
||||
NY = self.flow_field.FIELD_SHAPE[1]
|
||||
|
||||
# Add main cylinder (obstacle)
|
||||
center = (10 * L0, (NY - 1) / 2, 0)
|
||||
self.flow_field.add_cylinder(center, L0)
|
||||
|
||||
# Add sensors
|
||||
sensor_y_positions = [
|
||||
(NY - 1) / 2 + 2 * L0, # Above centerline
|
||||
(NY - 1) / 2, # At centerline
|
||||
(NY - 1) / 2 - 2 * L0, # Below centerline
|
||||
]
|
||||
for i in range(min(self.n_sensors, len(sensor_y_positions))):
|
||||
center = (40 * L0, sensor_y_positions[i], 0)
|
||||
self.flow_field.add_sensor(center, L0 / 4)
|
||||
|
||||
# Warmup simulation
|
||||
warmup_steps = int(self.warmup_steps_factor * NX / U0)
|
||||
self.flow_field.run(warmup_steps, np.zeros(self.n_control + 1, dtype=self.dtype))
|
||||
|
||||
# Collect baseline states for normalization
|
||||
for _ in range(self.fifo_length):
|
||||
self.flow_field.run(
|
||||
self.sample_interval,
|
||||
np.zeros(self.n_control + 1, dtype=self.dtype)
|
||||
)
|
||||
new_state = self.flow_field.obs.copy()[2:2 + self.n_sensors * 2]
|
||||
self.target_states = np.vstack((self.target_states, new_state))
|
||||
self.fifo_states.append(new_state)
|
||||
|
||||
# Calculate normalization factors
|
||||
self._calculate_normalization()
|
||||
|
||||
def _calculate_normalization(self):
|
||||
"""Calculate normalization factors from baseline states."""
|
||||
if len(self.target_states) > 0:
|
||||
self.sens_norm_fact = np.std(self.target_states, axis=0) + 1e-6
|
||||
self.sens_deviation = np.mean(self.target_states, axis=0)
|
||||
|
||||
def _normalize_state(self, state: np.ndarray) -> np.ndarray:
|
||||
"""Normalize state using calculated factors."""
|
||||
return (state - self.sens_deviation) / self.sens_norm_fact
|
||||
|
||||
def _compute_reward(self, state: np.ndarray, action: np.ndarray) -> float:
|
||||
"""
|
||||
Compute reward based on drag reduction and flow similarity.
|
||||
|
||||
Args:
|
||||
state: Current state observation
|
||||
action: Applied action
|
||||
|
||||
Returns:
|
||||
Total reward
|
||||
"""
|
||||
# Get force measurements from simulation
|
||||
obs = self.flow_field.obs
|
||||
cd = obs[0] # Drag coefficient
|
||||
cl = obs[1] # Lift coefficient
|
||||
|
||||
# Drag reduction reward (negative drag is good)
|
||||
self.reward_cd = -cd * 0.1
|
||||
|
||||
# Lift minimization (want symmetric flow)
|
||||
self.reward_cl = -abs(cl) * 0.05
|
||||
|
||||
# Flow similarity to baseline (want smooth control)
|
||||
if len(self.fifo_states) >= self.convergence_length:
|
||||
recent_states = np.array(list(self.fifo_states)[-self.convergence_length:])
|
||||
target_recent = self.target_states[-self.convergence_length:]
|
||||
|
||||
# Dynamic Time Warping distance (simplified)
|
||||
diff = np.mean(np.abs(recent_states - target_recent))
|
||||
self.reward_sim = -diff * 0.5
|
||||
else:
|
||||
self.reward_sim = 0.0
|
||||
|
||||
# Total reward
|
||||
total_reward = self.reward_cd + self.reward_cl + self.reward_sim
|
||||
|
||||
return float(total_reward)
|
||||
|
||||
def reset(
|
||||
self,
|
||||
seed: Optional[int] = None,
|
||||
options: Optional[Dict[str, Any]] = None
|
||||
) -> Tuple[np.ndarray, Dict[str, Any]]:
|
||||
"""
|
||||
Reset the environment to initial state.
|
||||
|
||||
Args:
|
||||
seed: Random seed for reproducibility
|
||||
options: Additional options
|
||||
|
||||
Returns:
|
||||
Tuple of (observation, info)
|
||||
"""
|
||||
super().reset(seed=seed)
|
||||
|
||||
self.current_step = 0
|
||||
self.fifo_states.clear()
|
||||
|
||||
# Run a few steps to get initial state
|
||||
for _ in range(10):
|
||||
self.flow_field.run(
|
||||
self.sample_interval,
|
||||
np.zeros(self.n_control + 1, dtype=self.dtype)
|
||||
)
|
||||
state = self.flow_field.obs.copy()[2:2 + self.n_sensors * 2]
|
||||
self.fifo_states.append(state)
|
||||
|
||||
# Get current state
|
||||
current_state = self.fifo_states[-1]
|
||||
|
||||
# Compute state derivative (approximation)
|
||||
if len(self.fifo_states) >= 2:
|
||||
state_derivative = self.fifo_states[-1] - self.fifo_states[-2]
|
||||
else:
|
||||
state_derivative = np.zeros_like(current_state)
|
||||
|
||||
# Normalize and concatenate
|
||||
obs = np.concatenate([
|
||||
self._normalize_state(current_state),
|
||||
self._normalize_state(state_derivative)
|
||||
]).astype(self.dtype)
|
||||
|
||||
info = {
|
||||
'step': self.current_step,
|
||||
'cd': self.flow_field.obs[0],
|
||||
'cl': self.flow_field.obs[1],
|
||||
}
|
||||
|
||||
return obs, info
|
||||
|
||||
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
|
||||
"""
|
||||
Take a step in the environment.
|
||||
|
||||
Args:
|
||||
action: Action to take (cylinder velocities)
|
||||
|
||||
Returns:
|
||||
Tuple of (observation, reward, terminated, truncated, info)
|
||||
"""
|
||||
# Convert action to control input
|
||||
# Action is in [-1, 1], scale to appropriate velocity range
|
||||
control = np.zeros(self.n_control + 1, dtype=self.dtype)
|
||||
control[:self.n_control] = action * 0.1 * self.config_field.velocity
|
||||
|
||||
# Run simulation
|
||||
self.flow_field.run(self.sample_interval, control)
|
||||
|
||||
# Get new state
|
||||
new_state = self.flow_field.obs.copy()[2:2 + self.n_sensors * 2]
|
||||
self.fifo_states.append(new_state)
|
||||
|
||||
# Compute observation
|
||||
if len(self.fifo_states) >= 2:
|
||||
state_derivative = self.fifo_states[-1] - self.fifo_states[-2]
|
||||
else:
|
||||
state_derivative = np.zeros_like(new_state)
|
||||
|
||||
obs = np.concatenate([
|
||||
self._normalize_state(new_state),
|
||||
self._normalize_state(state_derivative)
|
||||
]).astype(self.dtype)
|
||||
|
||||
# Compute reward
|
||||
reward = self._compute_reward(new_state, action)
|
||||
|
||||
# Check termination
|
||||
self.current_step += 1
|
||||
terminated = False # CFD simulations typically don't have natural termination
|
||||
truncated = self.current_step >= self.max_steps
|
||||
|
||||
# Info
|
||||
info = {
|
||||
'step': self.current_step,
|
||||
'cd': float(self.flow_field.obs[0]),
|
||||
'cl': float(self.flow_field.obs[1]),
|
||||
'reward_cd': float(self.reward_cd),
|
||||
'reward_cl': float(self.reward_cl),
|
||||
'reward_sim': float(self.reward_sim),
|
||||
}
|
||||
|
||||
return obs, reward, terminated, truncated, info
|
||||
|
||||
def render(self):
|
||||
"""Render the environment (not implemented)."""
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
"""Clean up resources."""
|
||||
if hasattr(self, 'flow_field'):
|
||||
del self.flow_field
|
||||
Reference in New Issue
Block a user