Frank14f/DynamisLab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

第一轮分析工作暂存

Browse Source
This commit is contained in:
Frank14f
2026-06-09 18:46:59 +08:00
parent 4612928398
commit d1b9922c6b
111 changed files with 24287 additions and 1581 deletions
Download Patch File Download Diff File Expand all files Collapse all files
archive/drl_pinball_new_api/scenes/__init__.py
+1
View File
@@ -0,0 +1 @@
archive/drl_pinball_new_api/scenes/karman_cloak/__init__.py
+1
View File
@@ -0,0 +1 @@
archive/drl_pinball_new_api/scenes/karman_cloak/re100_scene.py
+536
View File
@@ -0,0 +1,536 @@
# drl_pinball/scenes/karman_cloak/re100_scene.py
"""
Karman cloak re100 scene — inference orchestration for the Karman vortex street
cloaking scenario at code Reynolds number 100 (Re_D=50).
This scene exactly replicates the flow configuration of:
env_karman_cloak_standard.py + model d1a3o12_re100
Geometry (in lattice units, L0=20):
Disturbance cylinder: center=(200, CENTER_Y), radius=20
3 sensors: x=800, y=CENTER_Y + [-40, 0, 40], radius=5
Front pinball: center=(600, CENTER_Y), radius=10
Bottom pinball: center=(626, CENTER_Y-15), radius=10
Top pinball: center=(626, CENTER_Y+15), radius=10
Usage::
from scenes.karman_cloak.re100_scene import KarmanRe100Scene
scene = KarmanRe100Scene(device_id=0)
# Record target
scene.create_target_env()
scene.record_target("output/target.npz")
# Build full env with pinball
scene.create_full_env()
scene.collect_norm("output/norm.json")
scene.build_checkpoints("output/")
# Inference
scene.load_steady("output/checkpoint_steady.h5")
results = scene.run_controlled(model, n_steps=200)
scene.export_fields("output/fields/", step_indices=[0, 50, 100, 200])
"""
from __future__ import annotations
import json
import os
import sys
from collections import deque
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
# Add src directory to sys.path for package imports
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
_SRC = os.path.join(_REPO, "src")
if _SRC not in sys.path:
sys.path.insert(0, _SRC)
from drl_pinball.cfd.pinball_env import PinballEnv
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
U0 = 0.01
L0 = 20.0
SAMPLE_INTERVAL = 800
FIFO_LEN = 150
CONV_LEN = 30
S_DIM = 12
A_DIM = 3
ACTION_SCALE = 8.0
ACTION_BIAS = (0.0, -4.0, 4.0) # front, bottom, top
DATA_TYPE = np.float32
# Pinball config for new API
NEW_LBM_CONFIG = os.path.join(_REPO, "configs", "config_lbm_pinball.json")
class KarmanRe100Scene:
"""Karman cloak re100 scene manager."""
def __init__(self, device_id: int = 0, viscosity: float = 0.004):
self.device_id = device_id
self.viscosity = viscosity
self.env: Optional[PinballEnv] = None
# Body IDs (set during create_target_env / create_full_env)
self.dist_cyl_id: Optional[int] = None
self.sensor_ids: List[int] = []
self.front_cyl_id: Optional[int] = None
self.bottom_cyl_id: Optional[int] = None
self.top_cyl_id: Optional[int] = None
# Recorded data
self.target_states: Optional[np.ndarray] = None
self.norm_data: Optional[Dict] = None
self.save_states: Optional[np.ndarray] = None
def _center_y(self) -> float:
"""Return the center y of the domain (NY-1)/2 = 255.5."""
return 255.5 # (512 - 1) / 2
# -------------------------------------------------------------------
# Phase 1: Target recording (disturbance cylinder + sensors only)
# -------------------------------------------------------------------
def create_target_env(self):
"""Create flow field with disturbance cylinder + 3 sensors (no pinball)."""
self.env = PinballEnv(
lbm_config_path=NEW_LBM_CONFIG,
body_config_path=None,
device_id=self.device_id,
viscosity=self.viscosity,
velocity=U0,
)
cy = self._center_y()
# Disturbance cylinder (upstream)
self.dist_cyl_id = self.env.add_disturbance_cylinder(
center=(10.0 * L0, cy), radius=L0
)
# 3 sensors at x=40*L0
for y_off in [2.0, 0.0, -2.0]:
sid = self.env.add_sensor(
center=(40.0 * L0, cy + y_off * L0), radius=L0 / 4.0
)
self.sensor_ids.append(sid)
# Rebuild
self.env.reinitialize()
def record_target(self, out_dir: str) -> str:
"""Record target sensor signals (disturbance only, no pinball).
Saves to {out_dir}/target.npz and returns the path.
"""
if self.env is None:
raise RuntimeError("Call create_target_env() first")
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, "target.npz")
# Stabilize
stabilize_steps = int(4 * 1280 / U0)
self.env.run_and_read(stabilize_steps)
# Record target
target_list = []
for _ in range(FIFO_LEN):
result = self.env.run_and_read(SAMPLE_INTERVAL)
sens_flat = self.env.get_sensor_array(result["sensors"])
target_list.append(sens_flat)
self.target_states = np.array(target_list, dtype=DATA_TYPE)
np.savez(out_path, target_states=self.target_states)
return out_path
# -------------------------------------------------------------------
# Phase 2: Full env (add pinball, compute norm, checkpoint)
# -------------------------------------------------------------------
def create_full_env(self):
"""Add pinball cylinders to existing target env (or create from scratch)."""
if self.env is None:
self.create_target_env()
cy = self._center_y()
# Front cylinder
self.front_cyl_id = self.env.add_cylinder(
center=(30.0 * L0, cy), radius=L0 / 2.0
)
# Bottom cylinder
self.bottom_cyl_id = self.env.add_cylinder(
center=(31.3 * L0, cy - 0.75 * L0), radius=L0 / 2.0
)
# Top cylinder
self.top_cyl_id = self.env.add_cylinder(
center=(31.3 * L0, cy + 0.75 * L0), radius=L0 / 2.0
)
self.env.reinitialize()
def collect_norm(self, out_dir: str) -> Dict:
"""Compute normalisation factors from zero-action rollout.
Saves to {out_dir}/norm.json.
Returns the norm dict.
"""
if self.env is None:
raise RuntimeError("Call create_full_env() first")
os.makedirs(out_dir, exist_ok=True)
cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
# Stabilize with pinball
stabilize_steps = int(4 * 1280 / U0)
self.env.run_and_read(stabilize_steps)
# Snapshot the steady state
self.env.snapshot()
# Zero-action rollout for norm
fifo = deque(maxlen=FIFO_LEN)
for _ in range(FIFO_LEN):
result = self.env.run_and_read(SAMPLE_INTERVAL, read_fields=False)
# Forces in legacy order: dist, front, bottom, top
force_arr = []
force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
for cid in cylinder_ids:
force_arr.extend(result["forces"].get(cid, [0, 0]))
sensors_flat = self.env.get_sensor_array(result["sensors"])
obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
fifo.append(obs_flat)
# Compute norm from fifo
temp_states = np.array(fifo, dtype=DATA_TYPE)
force_norm_fact = 6.0 * float(np.max(np.abs(temp_states[:, 6:12])))
sens_deviation = np.mean(temp_states[:, 0:6], axis=0)
sens_norm_fact = np.zeros(6, dtype=DATA_TYPE)
for i in range(6):
sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp_states[:, i] - sens_deviation[i])))
self.norm_data = {
"force_norm_fact": force_norm_fact,
"sens_deviation": sens_deviation.tolist(),
"sens_norm_fact": sens_norm_fact.tolist(),
"action_bias": list(ACTION_BIAS),
"action_scale": ACTION_SCALE,
}
with open(os.path.join(out_dir, "norm.json"), "w") as f:
json.dump(self.norm_data, f, indent=2)
# Bias-action rollout for save_states
self.env.restore()
bias_omegas = {
self.front_cyl_id: float(ACTION_BIAS[0] * U0),
self.bottom_cyl_id: float(ACTION_BIAS[1] * U0),
self.top_cyl_id: float(ACTION_BIAS[2] * U0),
}
fifo.clear()
for _ in range(FIFO_LEN):
result = self.env.run_and_read(SAMPLE_INTERVAL, omegas=bias_omegas)
force_arr = []
force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
for cid in cylinder_ids:
force_arr.extend(result["forces"].get(cid, [0, 0]))
sensors_flat = self.env.get_sensor_array(result["sensors"])
obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
fifo.append(obs_flat)
self.save_states = np.array(fifo, dtype=DATA_TYPE)
np.savez(os.path.join(out_dir, "save_states.npz"), save_states=self.save_states)
# Save norm data as NPZ for easy loading
np.savez(
os.path.join(out_dir, "norm_data.npz"),
force_norm_fact=np.array([force_norm_fact], dtype=np.float32),
sens_deviation=np.array(sens_deviation, dtype=np.float32),
sens_norm_fact=np.array(sens_norm_fact, dtype=np.float32),
)
# Restore steady state
self.env.restore()
return self.norm_data
def build_checkpoints(self, out_dir: str):
"""Save steady-state and bias-state checkpoints."""
os.makedirs(out_dir, exist_ok=True)
# Steady state (already snapshot'd in collect_norm)
self.env.snapshot()
steady_path = os.path.join(out_dir, "checkpoint_steady.h5")
self.env.save_checkpoint(steady_path)
# Bias state: restore + run bias + save
self.env.restore()
cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
bias_omegas = {
self.front_cyl_id: float(ACTION_BIAS[0] * U0),
self.bottom_cyl_id: float(ACTION_BIAS[1] * U0),
self.top_cyl_id: float(ACTION_BIAS[2] * U0),
}
self.env.run_and_read(SAMPLE_INTERVAL * 10, omegas=bias_omegas)
bias_path = os.path.join(out_dir, "checkpoint_bias.h5")
self.env.save_checkpoint(bias_path)
# Restore steady
self.env.restore()
# -------------------------------------------------------------------
# Inference
# -------------------------------------------------------------------
def load_checkpoint(self, path: str):
"""Load from a saved checkpoint."""
if self.env is None:
raise RuntimeError("Env not created. Call create_full_env() first.")
self.env.load_checkpoint(path)
def run_uncontrolled(self, n_steps: int, out_dir: str) -> Dict:
"""Run uncontrolled inference (zero action)."""
os.makedirs(out_dir, exist_ok=True)
cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
sens_list, forc_list = [], []
for _ in range(n_steps):
result = self.env.run_and_read(SAMPLE_INTERVAL)
force_arr = []
force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
for cid in cylinder_ids:
force_arr.extend(result["forces"].get(cid, [0, 0]))
sensors_flat = self.env.get_sensor_array(result["sensors"])
sens_list.append(sensors_flat)
forc_list.append(np.array(force_arr, dtype=np.float32))
out = {
"sensors": np.array(sens_list, dtype=np.float32),
"forces": np.array(forc_list, dtype=np.float32),
}
np.savez(os.path.join(out_dir, "uncontrolled.npz"), **out)
# Final vorticity
self.env.export_vorticity_png(
os.path.join(out_dir, "vorticity_uncontrolled.png"),
title="Karman re100 uncontrolled",
)
return out
def run_controlled(
self,
model: Any,
n_steps: int,
out_dir: str,
*,
field_steps: Optional[List[int]] = None,
) -> Dict:
"""Run controlled inference with a PPO model.
Parameters
----------
model : PPO
Trained PPO model (must have Sin activation).
n_steps : int
Number of inference steps.
out_dir : str
Output directory.
field_steps : list of int, optional
Step indices at which to save Tecplot field files.
Returns
-------
dict with sensors, forces, obs, actions, rewards.
"""
os.makedirs(out_dir, exist_ok=True)
if field_steps is not None:
os.makedirs(os.path.join(out_dir, "fields"), exist_ok=True)
cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
cyl_map = {
self.front_cyl_id: 0,
self.bottom_cyl_id: 1,
self.top_cyl_id: 2,
}
norm = self.norm_data
if norm is None:
raise RuntimeError("Call collect_norm() first")
force_norm_fact = float(norm["force_norm_fact"])
sens_deviation = np.array(norm["sens_deviation"], dtype=np.float32)
sens_norm_fact = np.array(norm["sens_norm_fact"], dtype=np.float32)
# Restore steady state
self.env.restore()
# Bias FIFO
fifo = deque(maxlen=FIFO_LEN)
bias_omegas = {
self.front_cyl_id: float(ACTION_BIAS[0] * U0),
self.bottom_cyl_id: float(ACTION_BIAS[1] * U0),
self.top_cyl_id: float(ACTION_BIAS[2] * U0),
}
for _ in range(FIFO_LEN):
result = self.env.run_and_read(SAMPLE_INTERVAL, omegas=bias_omegas)
force_arr = []
force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
for cid in cylinder_ids:
force_arr.extend(result["forces"].get(cid, [0, 0]))
sensors_flat = self.env.get_sensor_array(result["sensors"])
obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
fifo.append(obs_flat)
sens_list, forc_list, obs_list = [], [], []
action_list, reward_list = [], []
reward_cd_list, reward_cl_list, reward_sim_list = [], [], []
obs = np.zeros(S_DIM, dtype=np.float32)
for step in range(n_steps):
# PPO action
action, _states = model.predict(obs, deterministic=True)
action = action.astype(np.float32).flatten()
action_list.append(action.copy())
# Convert to omegas
omegas = {}
for i, cid in enumerate(cylinder_ids):
omega_val = (action[i] * ACTION_SCALE + ACTION_BIAS[i]) * U0
omegas[cid] = float(omega_val)
# Run CFD
result = self.env.run_and_read(SAMPLE_INTERVAL, omegas=omegas)
# Build obs slice
force_arr = []
force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
for cid in cylinder_ids:
force_arr.extend(result["forces"].get(cid, [0, 0]))
sensors_flat = self.env.get_sensor_array(result["sensors"])
obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
fifo.append(obs_flat)
sens_list.append(sensors_flat)
forc_list.append(np.array(force_arr, dtype=np.float32))
# Build normalised observation
forces_norm = np.array(force_arr[2:], dtype=np.float32) / force_norm_fact # skip dist cy forces
sens_norm = (sensors_flat - sens_deviation) / sens_norm_fact
obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)
obs_list.append(obs)
# Compute reward
states_arr = np.array(fifo, dtype=np.float32)
if len(states_arr) >= CONV_LEN:
forces = states_arr[-1, 6:12] / force_norm_fact
cd = float((forces[0] + forces[2] + forces[4]) / 3.0)
cl = float((forces[1] + forces[3] + forces[5]) / 3.0)
sim = self._compute_similarity(states_arr)
r_cd = float(np.exp(-abs(cd * 20.0)))
r_cl = float(np.exp(-abs(cl * 80.0)))
r_sim = float(np.exp(-10.0 * abs(sim - 1.0)))
reward = float(min(0.3 * r_cd + 0.4 * r_cl + 0.3 * r_sim, 1.0))
else:
reward = 0.0
r_cd = r_cl = r_sim = 0.0
reward_list.append(reward)
reward_cd_list.append(r_cd)
reward_cl_list.append(r_cl)
reward_sim_list.append(r_sim)
# Field export
if field_steps is not None and step in field_steps:
fname = os.path.join(out_dir, "fields", f"field_{step:06d}.dat")
self.env.save_field_tecplot(fname)
out = {
"sensors": np.array(sens_list, dtype=np.float32),
"forces": np.array(forc_list, dtype=np.float32),
"obs": np.array(obs_list, dtype=np.float32),
"actions": np.array(action_list, dtype=np.float32),
"rewards": np.array(reward_list, dtype=np.float32),
"reward_cd": np.array(reward_cd_list, dtype=np.float32),
"reward_cl": np.array(reward_cl_list, dtype=np.float32),
"reward_sim": np.array(reward_sim_list, dtype=np.float32),
}
np.savez(os.path.join(out_dir, "controlled.npz"), **out)
# Final vorticity
self.env.export_vorticity_png(
os.path.join(out_dir, "vorticity_controlled.png"),
title="Karman re100 controlled (PPO)",
)
return out
def _compute_similarity(self, states_arr: np.ndarray) -> float:
"""Compute lag-compensated DTW similarity (matches legacy env logic)."""
if self.target_states is None:
return 0.0
target = self.target_states
# Lag from middle sensor (index 1 = sensor1_uy in sensor[6] block)
ref = target[CONV_LEN:2 * CONV_LEN, 1]
cur = states_arr[-CONV_LEN:, 1]
lag = self._calc_lag(ref, cur)
sim_sum = 0.0
for i in range(6):
t_seq = np.roll(target[:, i], -lag)[CONV_LEN:2 * CONV_LEN]
s_seq = states_arr[-CONV_LEN:, i]
sim_sum += self._calc_dtw_sim(t_seq, s_seq) / 6.0
return float(sim_sum)
@staticmethod
def _calc_lag(target: np.ndarray, state: np.ndarray) -> int:
tm = np.mean(target)
sm = np.mean(state)
corr = np.correlate(target - tm, state - sm, mode="full")
lags = np.arange(-len(target) + 1, len(target))
return int(lags[np.argmax(corr)])
@staticmethod
def _calc_dtw_sim(target: np.ndarray, state: np.ndarray) -> float:
n, m = len(target), len(state)
dtw = np.full((n + 1, m + 1), np.inf)
dtw[0, 0] = 0.0
for i in range(1, n + 1):
for j in range(1, m + 1):
cost = abs(float(target[i - 1]) - float(state[j - 1]))
dtw[i, j] = cost + min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1])
return float(1.0 - dtw[n, m] / n)
def close(self):
if self.env is not None:
self.env.close()
self.env = None