770 lines
28 KiB
Python
770 lines
28 KiB
Python
# CCD_analysis/scripts/phase1_collect.py
|
|
"""Phase 1: Data collection for all 4 analysis cases.
|
|
|
|
Usage::
|
|
conda run -n pycuda_3_10 python phase1_collect.py --case illusion --device 2
|
|
conda run -n pycuda_3_10 python phase1_collect.py --case cloak --device 3
|
|
conda run -n pycuda_3_10 python phase1_collect.py --case uncontrolled --device 3
|
|
conda run -n pycuda_3_10 python phase1_collect.py --case target_cylinder --device 2
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from collections import deque
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
# Add project root
|
|
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
|
if _REPO not in sys.path:
|
|
sys.path.insert(0, _REPO)
|
|
|
|
# Add analysis dir
|
|
_ANALYSIS = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
if _ANALYSIS not in sys.path:
|
|
sys.path.insert(0, _ANALYSIS)
|
|
|
|
from LegacyCelerisLab import FlowField # noqa: E402
|
|
|
|
from scripts.cfg import ( # noqa: E402
|
|
CONFIG_DIR, OUTPUT_DIR, U0, L0, NX, NY, CENTER_Y, DATA_TYPE,
|
|
PINBALL_RADIUS, FRONT_CENTER, BOTTOM_CENTER, TOP_CENTER,
|
|
ILLUSION_FRONT, ILLUSION_BOTTOM, ILLUSION_TOP,
|
|
SENSOR_RADIUS, SENSOR_CENTERS_CLOAK, SENSOR_CENTERS_ILLUSION,
|
|
TARGET_CYLINDER_CENTER, TARGET_CYLINDER_RADIUS,
|
|
SAMPLE_INTERVAL, SAMPLE_INTERVAL_ILLUSION,
|
|
ACTION_SCALE_CLOAK, ACTION_BIAS_CLOAK,
|
|
ACTION_SCALE_ILLUSION, ACTION_BIAS_ILLUSION,
|
|
MODEL_CLOAK_RE100, MODEL_ILLUSION_1L,
|
|
STABILIZE_STEPS, FIFO_LEN, N_PTS_PER_CYCLE,
|
|
nu_from_re,
|
|
)
|
|
from scripts.utils import ( # noqa: E402
|
|
load_configs, get_velocity_field, detect_cycle_stability,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PPO model loader (with Sin activation)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _load_ppo_model(model_path: str, device: str, s_dim: int = 12, a_dim: int = 3):
|
|
"""Load PPO model with Sin activation."""
|
|
import torch
|
|
from torch.nn import Module
|
|
from stable_baselines3 import PPO
|
|
import gymnasium as gym
|
|
from gymnasium import spaces
|
|
|
|
class Sin(Module):
|
|
def forward(self, x):
|
|
return torch.sin(x)
|
|
|
|
class DummyEnv(gym.Env):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.observation_space = spaces.Box(
|
|
low=-1, high=1, shape=(s_dim,), dtype=np.float32)
|
|
self.action_space = spaces.Box(
|
|
low=-1, high=1, shape=(a_dim,), dtype=np.float32)
|
|
def reset(self, seed=None):
|
|
return np.zeros(s_dim, dtype=np.float32), {}
|
|
def step(self, action):
|
|
return np.zeros(s_dim, dtype=np.float32), 0.0, False, False, {}
|
|
def render(self):
|
|
pass
|
|
|
|
dummy = DummyEnv()
|
|
model = PPO.load(model_path, env=dummy, device=device)
|
|
return model
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Field saving interval calculator
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _calc_save_interval(T_ref: float, n_pts_per_cycle: int = 24) -> int:
|
|
"""Calculate field save interval to get ~n_pts_per_cycle per cycle."""
|
|
interval = int(T_ref / n_pts_per_cycle)
|
|
return max(1, interval)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Phase 1a: Illusion
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def collect_illusion(device_id: int, data: dict) -> dict:
|
|
"""Collect illusion case data with proper norm computation and PPO inference.
|
|
|
|
Follows legacy_env_imit.py __init__ + step() logic exactly:
|
|
1. Target cylinder recording (separate FlowField)
|
|
2. FFT harmonics on target signals
|
|
3. Pinball env with norm computation
|
|
4. Bias-action FIFO initialization
|
|
5. PPO deterministic rollout with 14-dim normalized observations
|
|
"""
|
|
actual_U0 = 0.02 # model is 2U
|
|
viscosity = nu_from_re(100.0, u0=actual_U0)
|
|
sample_interval = SAMPLE_INTERVAL_ILLUSION # 600
|
|
fifo_len = 150
|
|
conv_len = 36
|
|
|
|
# ---- Step 1: Target cylinder recording ----
|
|
print("--- Record target cylinder ---")
|
|
target_U0 = actual_U0
|
|
target_nu = viscosity
|
|
cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
|
|
field_cfg = field_cfg._replace(viscosity=float(target_nu),
|
|
velocity=float(target_U0))
|
|
ff_target = FlowField(field_cfg, cuda_cfg, device_id=device_id)
|
|
|
|
# Target cylinder: center=(20*L0, CENTER_Y), radius=1.0*L0
|
|
L0 = 20.0
|
|
ff_target.add_cylinder(
|
|
(20.0 * L0, (512 - 1) / 2, 0.0), 1.0 * L0
|
|
)
|
|
# 3 sensors at x=30*L0
|
|
for y_off in [2.0, 0.0, -2.0]:
|
|
ff_target.add_sensor(
|
|
(30.0 * L0, (512 - 1) / 2 + y_off * L0, 0.0), L0 / 4.0
|
|
)
|
|
n_obj_target = ff_target.obs.size // 2 # 4
|
|
# Stabilize
|
|
ff_target.run(int(4 * 1280 / target_U0), np.zeros(n_obj_target, dtype=np.float32))
|
|
|
|
# Record 150 steps of obs[0:8] (3 sensors + 1 cylinder force)
|
|
target_states = np.empty((0, 8), dtype=np.float32)
|
|
for _ in range(fifo_len):
|
|
ff_target.run(sample_interval, np.zeros(n_obj_target, dtype=np.float32))
|
|
new_state = ff_target.obs.copy()[0:8]
|
|
target_states = np.vstack((target_states, new_state))
|
|
|
|
# FFT harmonics analysis
|
|
def analyze_harmonics(states, n_harmonics=5):
|
|
N, D = states.shape
|
|
result = []
|
|
for d in range(D):
|
|
y = states[:, d]
|
|
fft_coef = np.fft.rfft(y)
|
|
freqs = np.fft.rfftfreq(N, d=1)
|
|
amps = 2.0 * np.abs(fft_coef) / N
|
|
phases = np.angle(fft_coef)
|
|
idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
|
|
harmonics = {
|
|
'dc': float(np.real(fft_coef[0]) / N),
|
|
'amps': amps[idx].tolist(),
|
|
'freqs': freqs[idx].tolist(),
|
|
'phases': phases[idx].tolist(),
|
|
}
|
|
result.append(harmonics)
|
|
return result
|
|
|
|
target_harmonics = analyze_harmonics(target_states, n_harmonics=5)
|
|
del ff_target
|
|
print(f" target harmonics computed for {len(target_harmonics)} channels")
|
|
|
|
# ---- Step 2: Pinball env creation ----
|
|
print("--- Build pinball env ---")
|
|
ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
|
|
|
|
for y_off in [2.0, 0.0, -2.0]:
|
|
ff.add_sensor(
|
|
(30.0 * L0, (512 - 1) / 2 + y_off * L0, 0.0), L0 / 4.0
|
|
)
|
|
ff.add_cylinder((19.0 * L0, (512 - 1) / 2, 0.0), L0 / 2.0)
|
|
ff.add_cylinder((20.3 * L0, (512 - 1) / 2 + 0.75 * L0, 0.0), L0 / 2.0)
|
|
ff.add_cylinder((20.3 * L0, (512 - 1) / 2 - 0.75 * L0, 0.0), L0 / 2.0)
|
|
|
|
n_obj = ff.obs.size // 2 # 6
|
|
assert n_obj == 6, f"Expected 6 objects, got {n_obj}"
|
|
|
|
# Stabilize
|
|
ff.run(int(4 * 1280 / actual_U0), np.zeros(n_obj, dtype=np.float32))
|
|
ff.get_ddf()
|
|
ff.save_ddf() # checkpoint
|
|
|
|
# ---- Step 3: Norm computation (zero-action rollout) ----
|
|
print("--- Compute norm ---")
|
|
fifo = deque(maxlen=fifo_len)
|
|
for _ in range(fifo_len):
|
|
ff.run(sample_interval, np.zeros(n_obj, dtype=np.float32))
|
|
fifo.append(ff.obs.copy()[0:12])
|
|
|
|
temp_states = np.array(fifo, dtype=np.float32)
|
|
force_norm_fact = 6.0 * float(np.max(np.abs(temp_states[:, 6:12])))
|
|
sens_deviation = np.mean(temp_states[:, 0:6], axis=0).astype(np.float32)
|
|
sens_norm_fact = np.zeros(6, dtype=np.float32)
|
|
for i in range(6):
|
|
sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp_states[:, i] - sens_deviation[i])))
|
|
|
|
print(f" force_norm_fact={force_norm_fact:.6f}")
|
|
print(f" sens_deviation={sens_deviation}")
|
|
print(f" sens_norm_fact={sens_norm_fact}")
|
|
|
|
# ---- Step 4: Bias-action FIFO initialization ----
|
|
print("--- Bias-action FIFO init ---")
|
|
ff.apply_ddf()
|
|
# bias action from legacy env: [0, 0, 0, 0, -1*U0, 1*U0]
|
|
bias_arr = np.zeros(n_obj, dtype=np.float32)
|
|
bias_arr[4] = -1.0 * actual_U0 # bottom
|
|
bias_arr[5] = 1.0 * actual_U0 # top
|
|
|
|
fifo.clear()
|
|
for _ in range(fifo_len):
|
|
ff.run(sample_interval, bias_arr)
|
|
fifo.append(ff.obs.copy()[0:12])
|
|
|
|
save_states = list(fifo)
|
|
ff.apply_ddf() # restore checkpoint for reset
|
|
|
|
# ---- Step 5: PPO inference with adaptive sampling ----
|
|
print("--- PPO deterministic rollout (adaptive sampling) ---")
|
|
import torch
|
|
device_str = f"cuda:{device_id}" if torch.cuda.is_available() else "cpu"
|
|
model = _load_ppo_model(MODEL_ILLUSION_1L, device=device_str, s_dim=14, a_dim=3)
|
|
model.set_random_seed(19)
|
|
|
|
n_steps = 200
|
|
|
|
# Compute adaptive field sampling interval from expected period
|
|
# St = 0.267, D = 40, expected f = St * U0 / D
|
|
f_expected = 0.2667 * actual_U0 / 40.0
|
|
T_expected = int(1.0 / f_expected) if f_expected > 0 else 7500
|
|
field_interval = max(1, int(T_expected / N_PTS_PER_CYCLE))
|
|
print(f" T_expected={T_expected} steps, field_interval={field_interval} "
|
|
f"(~{T_expected/field_interval:.0f} pts/cycle)")
|
|
|
|
# Data at PPO-action cadence (once per 600 steps, for PPO state only)
|
|
ppo_actions = []
|
|
ppo_sensors_600 = []
|
|
|
|
# Dense data at field_interval cadence (for phase analysis)
|
|
dense_sensors = []
|
|
dense_forces = []
|
|
dense_ux = []
|
|
dense_uy = []
|
|
|
|
# Re-initialize FIFO for inference
|
|
fifo = deque(maxlen=fifo_len)
|
|
for state in save_states:
|
|
fifo.append(np.array(state, dtype=np.float32))
|
|
|
|
obs = np.zeros(14, dtype=np.float32)
|
|
|
|
for step in range(n_steps):
|
|
# PPO action
|
|
action, _states = model.predict(obs, deterministic=True)
|
|
action = action.astype(np.float32).flatten()
|
|
ppo_actions.append(action.copy())
|
|
|
|
# Convert to physical omega
|
|
temp = np.zeros(n_obj, dtype=np.float32)
|
|
omega = (action * ACTION_SCALE_ILLUSION
|
|
+ np.array(ACTION_BIAS_ILLUSION, dtype=np.float32)) * actual_U0
|
|
temp[3:6] = omega
|
|
|
|
# Run CFD with dense intra-step sampling
|
|
ff.context.push()
|
|
try:
|
|
# First chunk
|
|
ff.run(field_interval, temp)
|
|
ux, uy = get_velocity_field(ff, u0=actual_U0)
|
|
dense_ux.append(ux)
|
|
dense_uy.append(uy)
|
|
dense_sensors.append(ff.obs.copy()[0:6])
|
|
dense_forces.append(ff.obs.copy()[6:12])
|
|
|
|
# Second chunk (remaining)
|
|
remaining = sample_interval - field_interval
|
|
if remaining > 0:
|
|
ff.run(remaining, temp)
|
|
ux, uy = get_velocity_field(ff, u0=actual_U0)
|
|
dense_ux.append(ux)
|
|
dense_uy.append(uy)
|
|
dense_sensors.append(ff.obs.copy()[0:6])
|
|
dense_forces.append(ff.obs.copy()[6:12])
|
|
finally:
|
|
ff.context.pop()
|
|
|
|
# PPO state: use last obs_slice
|
|
last_sens = dense_sensors[-1]
|
|
last_force = dense_forces[-1]
|
|
obs_slice = np.concatenate([last_sens, last_force])
|
|
fifo.append(obs_slice)
|
|
ppo_sensors_600.append(obs_slice)
|
|
|
|
# Build normalized 14-dim observation for next PPO step
|
|
forces_norm = last_force / force_norm_fact
|
|
sens_norm = (last_sens - sens_deviation) / sens_norm_fact
|
|
target_recon = _gen_target_states_at(step, target_harmonics)
|
|
target_cd_norm = float(target_recon[0]) / force_norm_fact
|
|
target_cl_norm = float(target_recon[1]) / force_norm_fact
|
|
obs = np.clip(
|
|
np.hstack([forces_norm, sens_norm, target_cd_norm, target_cl_norm]),
|
|
-1.0, 1.0,
|
|
).astype(np.float32)
|
|
|
|
if step % 20 == 0:
|
|
print(f" step {step}/{n_steps}, action={action[0]:.3f} {action[1]:.3f} {action[2]:.3f}")
|
|
|
|
# Save dense data (for phase resampling)
|
|
ux_all = np.stack(dense_ux, axis=0)
|
|
uy_all = np.stack(dense_uy, axis=0)
|
|
dense_sensors_arr = np.array(dense_sensors, dtype=np.float32)
|
|
dense_forces_arr = np.array(dense_forces, dtype=np.float32)
|
|
ppo_actions_arr = np.array(ppo_actions, dtype=np.float32)
|
|
n_dense_per_step = len(dense_sensors) // n_steps
|
|
dense_dt = sample_interval / n_dense_per_step if n_dense_per_step > 0 else sample_interval
|
|
print(f" Dense sampling: {len(dense_sensors)} samples, "
|
|
f"{n_dense_per_step} per PPO step, dt={dense_dt:.0f} LBM steps")
|
|
|
|
out_dir = os.path.join(OUTPUT_DIR, "illusion")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
np.savez_compressed(os.path.join(out_dir, "fields.npz"), ux=ux_all, uy=uy_all)
|
|
np.savez(os.path.join(out_dir, "dense_sensors.npz"),
|
|
sensors=dense_sensors_arr, forces=dense_forces_arr,
|
|
dense_dt=dense_dt,
|
|
sample_interval=sample_interval)
|
|
|
|
# Save PPO-step-cadence data and metadata
|
|
np.savez(os.path.join(out_dir, "sensors.npz"),
|
|
sensors=dense_sensors_arr.reshape(n_steps, -1, 6)[:, -1],
|
|
forces=dense_forces_arr.reshape(n_steps, -1, 6)[:, -1],
|
|
actions=ppo_actions_arr,
|
|
sample_interval=sample_interval,
|
|
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))
|
|
|
|
# Save target data for later use
|
|
np.savez(os.path.join(out_dir, "target_harmonics.npz"),
|
|
target_states=target_states,
|
|
harmonics_data=np.array(target_harmonics, dtype=object))
|
|
|
|
meta = {
|
|
"case": "illusion",
|
|
"model": str(MODEL_ILLUSION_1L),
|
|
"n_steps": n_steps,
|
|
"n_fields": len(dense_ux),
|
|
"n_dense_samples": len(dense_sensors),
|
|
"dense_dt": dense_dt,
|
|
"T_expected": T_expected,
|
|
"field_interval": field_interval,
|
|
"sample_interval": sample_interval,
|
|
"action_scale": ACTION_SCALE_ILLUSION,
|
|
"action_bias": list(ACTION_BIAS_ILLUSION),
|
|
"U0": actual_U0,
|
|
"viscosity": viscosity,
|
|
"n_obj": n_obj,
|
|
"force_norm_fact": force_norm_fact,
|
|
"sens_deviation": sens_deviation.tolist(),
|
|
"sens_norm_fact": sens_norm_fact.tolist(),
|
|
}
|
|
with open(os.path.join(out_dir, "meta.json"), "w") as f:
|
|
json.dump(meta, f, indent=2)
|
|
|
|
print(f" Saved {len(dense_ux)} fields, {len(dense_sensors)} dense samples")
|
|
del ff, model
|
|
return meta
|
|
|
|
|
|
def _gen_target_states_at(t, harmonics):
|
|
"""Reconstruct target observable at step index t from harmonics.
|
|
|
|
Mirrors legacy_env_imit.py gen_target_states_at().
|
|
"""
|
|
t = np.asarray(t)
|
|
D = len(harmonics)
|
|
result = np.zeros((t.size, D), dtype=np.float32)
|
|
for d, h in enumerate(harmonics):
|
|
val = np.full(t.shape, h['dc'], dtype=np.float32)
|
|
amps = h['amps']
|
|
freqs = h['freqs']
|
|
phases = h['phases']
|
|
for amp, freq, phase in zip(amps, freqs, phases):
|
|
val += amp * np.cos(2 * np.pi * freq * t + phase)
|
|
result[:, d] = val
|
|
if result.shape[0] == 1:
|
|
return result[0]
|
|
return result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Phase 1b: Cloak (steady flow case)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def collect_cloak(device_id: int, data: dict) -> dict:
|
|
"""Collect cloak case data (PPO -> steady action -> mean flow)."""
|
|
viscosity = nu_from_re(100.0)
|
|
sample_interval = SAMPLE_INTERVAL
|
|
|
|
import torch
|
|
device_str = f"cuda:{device_id}" if torch.cuda.is_available() else "cpu"
|
|
model = _load_ppo_model(MODEL_CLOAK_RE100, device=device_str, s_dim=12, a_dim=3)
|
|
model.set_random_seed(0)
|
|
|
|
# Create env: 6 objects (3 sensors + 3 pinball, NO disturbance cylinder)
|
|
cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
|
|
field_cfg = field_cfg._replace(viscosity=float(viscosity))
|
|
ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
|
|
|
|
for sc in SENSOR_CENTERS_CLOAK:
|
|
ff.add_sensor((sc[0], sc[1], 0.0), SENSOR_RADIUS)
|
|
ff.add_cylinder((FRONT_CENTER[0], FRONT_CENTER[1], 0.0), PINBALL_RADIUS)
|
|
ff.add_cylinder((BOTTOM_CENTER[0], BOTTOM_CENTER[1], 0.0), PINBALL_RADIUS)
|
|
ff.add_cylinder((TOP_CENTER[0], TOP_CENTER[1], 0.0), PINBALL_RADIUS)
|
|
|
|
n_obj = ff.obs.size // 2
|
|
assert n_obj == 6, f"Expected 6 objects for cloak, got {n_obj}"
|
|
|
|
# Stabilize
|
|
ff.run(STABILIZE_STEPS, np.zeros(n_obj, dtype=np.float32))
|
|
|
|
# ---- PPO deterministic rollout to find steady action ----
|
|
n_ppo_steps = 200
|
|
print(f"Running {n_ppo_steps} PPO steps to extract steady action...")
|
|
|
|
obs = np.zeros(12, dtype=np.float32)
|
|
actions_list = []
|
|
sensors_list = []
|
|
forces_list = []
|
|
|
|
for step in range(n_ppo_steps):
|
|
action, _states = model.predict(obs, deterministic=True)
|
|
action = action.astype(np.float32).flatten()
|
|
actions_list.append(action.copy())
|
|
|
|
temp = np.zeros(n_obj, dtype=np.float32)
|
|
omega = (action * ACTION_SCALE_CLOAK
|
|
+ np.array(ACTION_BIAS_CLOAK, dtype=np.float32)) * U0
|
|
temp[3:6] = omega
|
|
|
|
ff.context.push()
|
|
try:
|
|
ff.run(sample_interval, temp)
|
|
finally:
|
|
ff.context.pop()
|
|
|
|
obs_slice = ff.obs.copy()[0:12]
|
|
sensors_list.append(obs_slice[0:6].copy())
|
|
forces_list.append(obs_slice[6:12].copy())
|
|
|
|
# Build observation for next step
|
|
obs = np.clip(np.hstack([obs_slice[6:12], obs_slice[0:6]]),
|
|
-10.0, 10.0).astype(np.float32)
|
|
|
|
# Extract steady action (average of last 100 steps)
|
|
actions_arr = np.array(actions_list, dtype=np.float32)
|
|
steady_action = np.mean(actions_arr[-100:], axis=0)
|
|
print(f" Steady action ([-1,1]): {steady_action[0]:.4f} {steady_action[1]:.4f} {steady_action[2]:.4f}")
|
|
print(f" Steady omega (U0 multiples): "
|
|
f"{(steady_action*ACTION_SCALE_CLOAK+np.array(ACTION_BIAS_CLOAK))[0]:.4f} "
|
|
f"{(steady_action*ACTION_SCALE_CLOAK+np.array(ACTION_BIAS_CLOAK))[1]:.4f} "
|
|
f"{(steady_action*ACTION_SCALE_CLOAK+np.array(ACTION_BIAS_CLOAK))[2]:.4f}")
|
|
|
|
# ---- Apply steady action and record mean flow ----
|
|
print("Applying steady action and recording...")
|
|
temp_steady = np.zeros(n_obj, dtype=np.float32)
|
|
omega_steady = (steady_action * ACTION_SCALE_CLOAK
|
|
+ np.array(ACTION_BIAS_CLOAK, dtype=np.float32)) * U0
|
|
temp_steady[3:6] = omega_steady
|
|
|
|
# Re-stabilize with steady action (4x NX/U0)
|
|
ff.context.push()
|
|
try:
|
|
ff.run(STABILIZE_STEPS, temp_steady)
|
|
finally:
|
|
ff.context.pop()
|
|
|
|
# Record steady state fields and sensors
|
|
n_steady_samples = 30
|
|
steady_sensors = []
|
|
steady_forces = []
|
|
steady_ux = []
|
|
steady_uy = []
|
|
|
|
for i in range(n_steady_samples):
|
|
ff.context.push()
|
|
try:
|
|
ff.run(sample_interval, temp_steady)
|
|
finally:
|
|
ff.context.pop()
|
|
obs_slice = ff.obs.copy()[0:12]
|
|
steady_sensors.append(obs_slice[0:6])
|
|
steady_forces.append(obs_slice[6:12])
|
|
ux, uy = get_velocity_field(ff, u0=U0)
|
|
steady_ux.append(ux)
|
|
steady_uy.append(uy)
|
|
|
|
steady_sensors_arr = np.array(steady_sensors, dtype=np.float32)
|
|
steady_forces_arr = np.array(steady_forces, dtype=np.float32)
|
|
ux_all = np.stack(steady_ux, axis=0)
|
|
uy_all = np.stack(steady_uy, axis=0)
|
|
|
|
out_dir = os.path.join(OUTPUT_DIR, "cloak")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
np.savez_compressed(os.path.join(out_dir, "fields.npz"),
|
|
ux=ux_all, uy=uy_all)
|
|
np.savez(os.path.join(out_dir, "sensors.npz"),
|
|
sensors=steady_sensors_arr, forces=steady_forces_arr)
|
|
np.savez(os.path.join(out_dir, "ppo_rollout.npz"),
|
|
actions=actions_arr,
|
|
sensors=np.array(sensors_list, dtype=np.float32),
|
|
forces=np.array(forces_list, dtype=np.float32),
|
|
steady_action=steady_action)
|
|
|
|
meta = {
|
|
"case": "cloak",
|
|
"model": str(MODEL_CLOAK_RE100),
|
|
"sample_interval": sample_interval,
|
|
"action_scale": ACTION_SCALE_CLOAK,
|
|
"action_bias": list(ACTION_BIAS_CLOAK),
|
|
"steady_action_norm": steady_action.tolist(),
|
|
"steady_omega_U0": (steady_action * ACTION_SCALE_CLOAK
|
|
+ np.array(ACTION_BIAS_CLOAK)).tolist(),
|
|
"U0": U0,
|
|
"viscosity": viscosity,
|
|
"n_obj": n_obj,
|
|
"n_steady_samples": n_steady_samples,
|
|
}
|
|
with open(os.path.join(out_dir, "meta.json"), "w") as f:
|
|
json.dump(meta, f, indent=2)
|
|
|
|
print(f" Steady action recorded. Mean sensors: "
|
|
f"{np.mean(steady_sensors_arr, axis=0)}")
|
|
print(f" Mean total force: "
|
|
f"Fx={np.mean(steady_forces_arr[:, 0::2]):.6f} "
|
|
f"Fy={np.mean(steady_forces_arr[:, 1::2]):.6f}")
|
|
del ff, model
|
|
return meta
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Phase 1c: Uncontrolled
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def collect_uncontrolled(device_id: int, data: dict) -> dict:
|
|
"""Collect uncontrolled case data (zero-action baseline)."""
|
|
viscosity = nu_from_re(100.0)
|
|
sample_interval = SAMPLE_INTERVAL
|
|
T_ref = data.get("T_ref", 15000.0)
|
|
save_interval = _calc_save_interval(T_ref)
|
|
|
|
cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
|
|
field_cfg = field_cfg._replace(viscosity=float(viscosity))
|
|
ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
|
|
|
|
for sc in SENSOR_CENTERS_CLOAK:
|
|
ff.add_sensor((sc[0], sc[1], 0.0), SENSOR_RADIUS)
|
|
ff.add_cylinder((FRONT_CENTER[0], FRONT_CENTER[1], 0.0), PINBALL_RADIUS)
|
|
ff.add_cylinder((BOTTOM_CENTER[0], BOTTOM_CENTER[1], 0.0), PINBALL_RADIUS)
|
|
ff.add_cylinder((TOP_CENTER[0], TOP_CENTER[1], 0.0), PINBALL_RADIUS)
|
|
|
|
n_obj = ff.obs.size // 2
|
|
assert n_obj == 6
|
|
|
|
# Stabilize
|
|
ff.run(STABILIZE_STEPS, np.zeros(n_obj, dtype=np.float32))
|
|
|
|
# Run uncontrolled
|
|
n_steps = 200
|
|
sensors_list = []
|
|
forces_list = []
|
|
ux_fields = []
|
|
uy_fields = []
|
|
|
|
for step in range(n_steps):
|
|
ff.context.push()
|
|
try:
|
|
remaining = sample_interval
|
|
while remaining > 0:
|
|
chunk = min(remaining, save_interval)
|
|
ff.run(chunk, np.zeros(n_obj, dtype=np.float32))
|
|
remaining -= chunk
|
|
ux, uy = get_velocity_field(ff, u0=U0)
|
|
ux_fields.append(ux)
|
|
uy_fields.append(uy)
|
|
finally:
|
|
ff.context.pop()
|
|
|
|
obs_slice = ff.obs.copy()[0:12]
|
|
sensors_list.append(obs_slice[0:6])
|
|
forces_list.append(obs_slice[6:12])
|
|
|
|
sensors = np.array(sensors_list, dtype=np.float32)
|
|
forces = np.array(forces_list, dtype=np.float32)
|
|
ux_all = np.stack(ux_fields, axis=0)
|
|
uy_all = np.stack(uy_fields, axis=0)
|
|
|
|
out_dir = os.path.join(OUTPUT_DIR, "uncontrolled")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
np.savez_compressed(os.path.join(out_dir, "fields.npz"),
|
|
ux=ux_all, uy=uy_all)
|
|
np.savez(os.path.join(out_dir, "sensors.npz"),
|
|
sensors=sensors, forces=forces)
|
|
|
|
meta = {
|
|
"case": "uncontrolled",
|
|
"U0": U0,
|
|
"viscosity": viscosity,
|
|
"n_steps": n_steps,
|
|
"n_fields": len(ux_fields),
|
|
"sample_interval": sample_interval,
|
|
"n_obj": n_obj,
|
|
}
|
|
with open(os.path.join(out_dir, "meta.json"), "w") as f:
|
|
json.dump(meta, f, indent=2)
|
|
|
|
print(f" Saved {len(ux_fields)} fields, {len(sensors)} sensor steps")
|
|
del ff
|
|
return meta
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Phase 1d: Target cylinder (reference for period detection)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def collect_target_cylinder(device_id: int, data: dict) -> dict:
|
|
"""Collect target 2D cylinder reference data.
|
|
|
|
Most data was already collected in Phase 0. Here we just ensure
|
|
the fields are properly saved with the right naming.
|
|
"""
|
|
# Phase 0 already saved data to output/target_cylinder/
|
|
# Just verify it exists and copy meta
|
|
out_dir = os.path.join(OUTPUT_DIR, "target_cylinder")
|
|
meta_path = os.path.join(out_dir, "meta.json")
|
|
if not os.path.exists(meta_path):
|
|
raise RuntimeError(
|
|
"Phase 0 must be run first. No target_cylinder data found."
|
|
)
|
|
|
|
with open(meta_path, "r") as f:
|
|
meta = json.load(f)
|
|
|
|
print(f"Target cylinder data found at {out_dir}")
|
|
print(f" f_ref={meta['f_ref']:.6f}, T_ref={meta['T_ref']:.0f}, St={meta['St']:.4f}")
|
|
print(f" CV_T={meta['CV_T']:.4f}")
|
|
return meta
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Empty channel (target steady flow for cloak comparison)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def collect_empty_channel(device_id: int) -> dict:
|
|
"""Run empty channel (no bodies) and record steady parabolic flow."""
|
|
viscosity = nu_from_re(100.0)
|
|
|
|
cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
|
|
field_cfg = field_cfg._replace(viscosity=float(viscosity))
|
|
ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
|
|
|
|
# Need at least one sensor (legacy API requirement)
|
|
ff.add_sensor((NX - 10, CENTER_Y, 0.0), SENSOR_RADIUS)
|
|
n_obj = ff.obs.size // 2
|
|
ff.run(STABILIZE_STEPS, np.zeros(n_obj, dtype=np.float32))
|
|
|
|
# Record a few fields
|
|
ux_list, uy_list = [], []
|
|
for i in range(5):
|
|
ff.run(SAMPLE_INTERVAL, np.zeros(n_obj, dtype=np.float32))
|
|
ux, uy = get_velocity_field(ff, u0=U0)
|
|
ux_list.append(ux)
|
|
uy_list.append(uy)
|
|
|
|
ux_all = np.stack(ux_list, axis=0)
|
|
uy_all = np.stack(uy_list, axis=0)
|
|
|
|
out_dir = os.path.join(OUTPUT_DIR, "empty_channel")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
np.savez_compressed(os.path.join(out_dir, "fields.npz"),
|
|
ux=ux_all, uy=uy_all)
|
|
|
|
meta = {
|
|
"case": "empty_channel",
|
|
"U0": U0,
|
|
"viscosity": viscosity,
|
|
"n_fields": len(ux_list),
|
|
}
|
|
with open(os.path.join(out_dir, "meta.json"), "w") as f:
|
|
json.dump(meta, f, indent=2)
|
|
|
|
print("Empty channel flow recorded")
|
|
del ff
|
|
return meta
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Phase 1: Data collection")
|
|
ap.add_argument("--case", type=str, required=True,
|
|
choices=["all", "illusion", "cloak", "uncontrolled",
|
|
"target_cylinder", "empty_channel"],
|
|
help="Case to collect")
|
|
ap.add_argument("--device", type=int, default=2, help="GPU device ID")
|
|
args = ap.parse_args()
|
|
|
|
# Load Phase 0 data for f_ref / T_ref
|
|
f_ref_path = os.path.join(OUTPUT_DIR, "target_cylinder", "meta.json")
|
|
if os.path.exists(f_ref_path):
|
|
with open(f_ref_path, "r") as f:
|
|
phase0_data = json.load(f)
|
|
else:
|
|
phase0_data = {"T_ref": 15000.0, "f_ref": 6.67e-5}
|
|
print("WARNING: Phase 0 not found, using default T_ref=15000")
|
|
|
|
t0 = time.time()
|
|
results = {}
|
|
|
|
if args.case in ("all", "illusion"):
|
|
print("=" * 60)
|
|
print("Collecting Illusion case...")
|
|
print("=" * 60)
|
|
phase0_data["illusion_2u"] = True
|
|
results["illusion"] = collect_illusion(args.device, phase0_data)
|
|
|
|
if args.case in ("all", "cloak"):
|
|
print("=" * 60)
|
|
print("Collecting Cloak case...")
|
|
print("=" * 60)
|
|
results["cloak"] = collect_cloak(args.device, phase0_data)
|
|
|
|
if args.case in ("all", "uncontrolled"):
|
|
print("=" * 60)
|
|
print("Collecting Uncontrolled case...")
|
|
print("=" * 60)
|
|
results["uncontrolled"] = collect_uncontrolled(args.device, phase0_data)
|
|
|
|
if args.case in ("all", "target_cylinder"):
|
|
print("=" * 60)
|
|
print("Collecting Target Cylinder...")
|
|
print("=" * 60)
|
|
results["target_cylinder"] = collect_target_cylinder(
|
|
args.device, phase0_data)
|
|
|
|
if args.case in ("all", "empty_channel"):
|
|
print("=" * 60)
|
|
print("Collecting Empty Channel (steady target)...")
|
|
print("=" * 60)
|
|
results["empty_channel"] = collect_empty_channel(args.device)
|
|
|
|
elapsed = time.time() - t0
|
|
print(f"\nPhase 1 complete in {elapsed:.1f}s")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|