329 lines
11 KiB
Python
329 lines
11 KiB
Python
"""
|
||
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
|