DynamisLab/src/CCD_analysis/scripts/phase1_collect.py
2026-06-09 18:46:59 +08:00

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())