第一轮分析工作暂存
This commit is contained in:
Frank14f
@@ -0,0 +1,371 @@
|
||||
# drl_pinball/cfd/pinball_env.py
|
||||
"""
|
||||
PinballEnv — wraps CelerisLab.Simulation for DRL inference.
|
||||
|
||||
This class provides the same telemetry interface as LegacyCelerisLab.FlowField.run(),
|
||||
but using the new Simulation API. The key difference is that the new API returns
|
||||
N-step cumulative values, while the old API returned per-step averages.
|
||||
|
||||
Usage::
|
||||
|
||||
from pinball_env import PinballEnv
|
||||
|
||||
env = PinballEnv(lbm_config, body_config, device_id=0)
|
||||
env.set_cylinders({front_id: 0.0, bottom_id: -0.04, top_id: 0.04})
|
||||
result = env.run_and_read(800)
|
||||
# result['forces'][body_id] = [fx_per_step, fy_per_step]
|
||||
# result['sensors'][body_id] = [ux_per_step, uy_per_step]
|
||||
|
||||
env.snapshot()
|
||||
env.restore()
|
||||
|
||||
env.save_field_tecplot("output.dat")
|
||||
env.export_vorticity_png("vorticity.png")
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from typing import Dict, List, Optional, Tuple, Any
|
||||
|
||||
import numpy as np
|
||||
import pycuda.driver as cuda
|
||||
|
||||
_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)
|
||||
_DEFAULT_LBM = os.path.join(_REPO, "configs", "config_lbm_pinball.json")
|
||||
_DEFAULT_BODY = os.path.join(_REPO, "configs", "config_body.json")
|
||||
|
||||
# LBM constants
|
||||
_CS2 = 1.0 / 3.0 # lattice speed of sound squared
|
||||
|
||||
|
||||
class PinballEnv:
|
||||
"""High-level wrapper around CelerisLab.Simulation for pinball DRL tasks.
|
||||
|
||||
Responsibilities:
|
||||
- Create and manage a Simulation instance
|
||||
- Provide run_and_read() that matches old API semantics (per-step averages)
|
||||
- Manage body ids for sensors and cylinders
|
||||
- Support snapshot/restore for checkpointing
|
||||
- Export macroscopic fields
|
||||
|
||||
Body ID convention (all envs follow this order):
|
||||
sensors[0], sensors[1], sensors[2], [disturbance_cylinder],
|
||||
front_cylinder, bottom_cylinder, top_cylinder
|
||||
|
||||
Some scenes (illusion, vortex) omit the disturbance cylinder.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
lbm_config_path: Optional[str] = None,
|
||||
body_config_path: Optional[str] = None,
|
||||
device_id: int = 0,
|
||||
*,
|
||||
viscosity: Optional[float] = None,
|
||||
velocity: Optional[float] = None,
|
||||
):
|
||||
# Build config with optional physics override
|
||||
if lbm_config_path is None:
|
||||
lbm_config_path = _DEFAULT_LBM
|
||||
if body_config_path is None:
|
||||
body_config_path = _DEFAULT_BODY
|
||||
|
||||
if viscosity is not None or velocity is not None:
|
||||
# Create a temp config with overridden physics
|
||||
with open(lbm_config_path) as f:
|
||||
cfg = json.load(f)
|
||||
if viscosity is not None:
|
||||
cfg["physics"]["viscosity"] = float(viscosity)
|
||||
if velocity is not None:
|
||||
cfg["physics"]["velocity"] = float(velocity)
|
||||
tmpd = tempfile.mkdtemp(prefix="pinball_env_cfg_")
|
||||
tmp_cfg_path = os.path.join(tmpd, "config_lbm.json")
|
||||
with open(tmp_cfg_path, "w") as f:
|
||||
json.dump(cfg, f, indent=2)
|
||||
lbm_config_path = tmp_cfg_path
|
||||
|
||||
from CelerisLab import Simulation
|
||||
|
||||
self.sim = Simulation(
|
||||
lbm_config_path=lbm_config_path,
|
||||
body_config_path=body_config_path,
|
||||
device_id=device_id,
|
||||
)
|
||||
|
||||
self._velocity = float(velocity) if velocity is not None else 0.01
|
||||
self._device_id = device_id
|
||||
self._stream = cuda.Stream()
|
||||
|
||||
# Body tracking
|
||||
self._body_ids: Dict[str, List[int]] = {
|
||||
"sensors": [],
|
||||
"cylinders": [],
|
||||
"disturbance": [],
|
||||
}
|
||||
self._body_id_to_name: Dict[int, str] = {}
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Geometry construction
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def add_cylinder(self, center: Tuple[float, float], radius: float) -> int:
|
||||
"""Add a cylinder body. Returns body_id."""
|
||||
from CelerisLab import Simulation
|
||||
body_id = self.sim.add_body("circle", center=center, radius=radius)
|
||||
self._body_ids["cylinders"].append(body_id)
|
||||
self._body_id_to_name[body_id] = f"cylinder_{len(self._body_ids['cylinders'])}"
|
||||
return body_id
|
||||
|
||||
def add_sensor(self, center: Tuple[float, float], radius: float) -> int:
|
||||
"""Add a sensor body. Returns body_id."""
|
||||
body_id = self.sim.add_body("sensor", center=center, radius=radius)
|
||||
self._body_ids["sensors"].append(body_id)
|
||||
self._body_id_to_name[body_id] = f"sensor_{len(self._body_ids['sensors'])}"
|
||||
return body_id
|
||||
|
||||
def add_disturbance_cylinder(self, center: Tuple[float, float], radius: float) -> int:
|
||||
"""Add a disturbance cylinder (upstream). Returns body_id."""
|
||||
body_id = self.sim.add_body("circle", center=center, radius=radius)
|
||||
self._body_ids["disturbance"].append(body_id)
|
||||
self._body_id_to_name[body_id] = "disturbance"
|
||||
return body_id
|
||||
|
||||
def reinitialize(self):
|
||||
"""Recompile and reinitialize after adding bodies."""
|
||||
self.sim.initialize()
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Runtime control
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def set_cylinder_omega(self, body_id: int, omega: float):
|
||||
"""Set cylinder rotation speed in lattice units."""
|
||||
self.sim.set_body(body_id, omega=float(omega))
|
||||
|
||||
def set_cylinders(self, omegas: Dict[int, float]):
|
||||
"""Set multiple cylinder omegas at once. {body_id: omega}."""
|
||||
for bid, omega in omegas.items():
|
||||
self.sim.set_body(bid, omega=float(omega))
|
||||
|
||||
def run_and_read(
|
||||
self,
|
||||
steps: int,
|
||||
omegas: Optional[Dict[int, float]] = None,
|
||||
*,
|
||||
read_fields: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run N LBM steps and read telemetry.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
steps : int
|
||||
Number of LBM steps to run.
|
||||
omegas : dict, optional
|
||||
Cylinder omegas to set before running. {body_id: omega}
|
||||
read_fields : bool
|
||||
If True, also return macroscopic field.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict with:
|
||||
forces : dict {body_id: [fx, fy]} per-step average
|
||||
sensors : dict {body_id: [ux, uy]} per-step average
|
||||
fields : dict (only if read_fields=True)
|
||||
"""
|
||||
# Set omegas if provided
|
||||
if omegas is not None:
|
||||
self.set_cylinders(omegas)
|
||||
|
||||
# Zero GPU telemetry
|
||||
self.sim.bodies.zero_force_segment_async(self._stream)
|
||||
if self.sim.field.n_sensor > 0:
|
||||
self.sim.bodies.zero_sensor_segment_async(self._stream)
|
||||
|
||||
# Run steps
|
||||
self.sim.stepper.step(
|
||||
int(steps),
|
||||
action_gpu=self.sim.bodies.action_gpu,
|
||||
obs_gpu=self.sim.bodies.obs_gpu,
|
||||
stream=self._stream,
|
||||
)
|
||||
|
||||
# Download telemetry
|
||||
self.sim.bodies.download_obs_full_async(self._stream)
|
||||
self._stream.synchronize()
|
||||
|
||||
# Read forces (per-step average)
|
||||
forces = {}
|
||||
for bid in self._body_ids["cylinders"]:
|
||||
f = self.sim.read_force(bid)
|
||||
forces[bid] = (np.array(f, dtype=np.float32) / float(steps)).tolist()
|
||||
|
||||
for bid in self._body_ids["disturbance"]:
|
||||
f = self.sim.read_force(bid)
|
||||
forces[bid] = (np.array(f, dtype=np.float32) / float(steps)).tolist()
|
||||
|
||||
# Read sensors (per-step average, raw sum / steps, NO cell count division)
|
||||
sensors = {}
|
||||
for bid in self._body_ids["sensors"]:
|
||||
s = self.sim.read_sensor(bid, normalize=False)
|
||||
sensors[bid] = (np.array(s, dtype=np.float32) / float(steps)).tolist()
|
||||
|
||||
result = {
|
||||
"forces": forces,
|
||||
"sensors": sensors,
|
||||
"n_steps": int(steps),
|
||||
}
|
||||
|
||||
if read_fields:
|
||||
macro = self.sim.get_macroscopic()
|
||||
result["fields"] = {
|
||||
"ux": np.asarray(macro["ux"], dtype=np.float32),
|
||||
"uy": np.asarray(macro["uy"], dtype=np.float32),
|
||||
"rho": np.asarray(macro["rho"], dtype=np.float32),
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
def get_sensor_array(self, sensors: Dict[int, list]) -> np.ndarray:
|
||||
"""Convert sensor dict to flat array in body_id order: [s0_ux, s0_uy, s1_ux, ...]."""
|
||||
arr = []
|
||||
for bid in sorted(sensors.keys()):
|
||||
arr.extend(sensors[bid])
|
||||
return np.array(arr, dtype=np.float32)
|
||||
|
||||
def get_force_array(self, forces: Dict[int, list]) -> np.ndarray:
|
||||
"""Convert force dict to flat array in cylinder body_id order."""
|
||||
arr = []
|
||||
for bid in self._body_ids["cylinders"]:
|
||||
if bid in forces:
|
||||
arr.extend(forces[bid])
|
||||
for bid in self._body_ids["disturbance"]:
|
||||
if bid in forces:
|
||||
arr.extend(forces[bid])
|
||||
return np.array(arr, dtype=np.float32)
|
||||
|
||||
def get_force_array_legacy_order(self, forces: Dict[int, list]) -> np.ndarray:
|
||||
"""Return forces in legacy obs order: dist_cyl first, then front, bottom, top.
|
||||
|
||||
This matches the old API's flat obs array layout:
|
||||
[dist_fx, dist_fy, front_fx, front_fy, bottom_fx, bottom_fy, top_fx, top_fy]
|
||||
"""
|
||||
arr = []
|
||||
# Disturbance cylinders first
|
||||
for bid in self._body_ids["disturbance"]:
|
||||
if bid in forces:
|
||||
arr.extend(forces[bid])
|
||||
# Then regular cylinders
|
||||
for bid in self._body_ids["cylinders"]:
|
||||
if bid in forces:
|
||||
arr.extend(forces[bid])
|
||||
return np.array(arr, dtype=np.float32)
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Checkpoint / Snapshot
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def snapshot(self):
|
||||
"""Save in-memory snapshot of current DDF state."""
|
||||
self.sim.snapshot()
|
||||
|
||||
def restore(self):
|
||||
"""Restore from in-memory snapshot."""
|
||||
self.sim.restore()
|
||||
|
||||
def save_checkpoint(self, path: str):
|
||||
"""Save HDF5 checkpoint to disk."""
|
||||
self.sim.save_checkpoint(path)
|
||||
|
||||
def load_checkpoint(self, path: str):
|
||||
"""Load HDF5 checkpoint from disk."""
|
||||
self.sim.load_checkpoint(path)
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Field export
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def get_macroscopic(self) -> Dict[str, np.ndarray]:
|
||||
"""Download macroscopic field. Returns {rho, ux, uy}."""
|
||||
return self.sim.get_macroscopic()
|
||||
|
||||
def save_field_tecplot(self, filename: str):
|
||||
"""Save current flow field in Tecplot format.
|
||||
|
||||
Matches the format of old save_field() in legacy envs.
|
||||
"""
|
||||
macro = self.get_macroscopic()
|
||||
ux = np.asarray(macro["ux"], dtype=np.float32)
|
||||
uy = np.asarray(macro["uy"], dtype=np.float32)
|
||||
|
||||
nx, ny = ux.shape
|
||||
u0 = self._velocity
|
||||
|
||||
with open(filename, "w") as f:
|
||||
f.write('Title= "LBM 2D"\r\n')
|
||||
f.write('VARIABLES= "X","Y","flag","U","V",\r\n')
|
||||
f.write(f"ZONE T= \"BOX\",I= {nx},J= {ny},F=POINT\r\n")
|
||||
for j in range(ny):
|
||||
for i in range(nx):
|
||||
u_val = ux[i, j] / u0
|
||||
v_val = uy[i, j] / u0
|
||||
f.write(f"{i},{j},0,{u_val},{v_val}\r\n")
|
||||
|
||||
def export_vorticity_png(self, path: str, title: str = ""):
|
||||
"""Export vorticity field as PNG."""
|
||||
try:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
macro = self.get_macroscopic()
|
||||
ux = np.asarray(macro["ux"], dtype=np.float64)
|
||||
uy = np.asarray(macro["uy"], dtype=np.float64)
|
||||
|
||||
omega = np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
|
||||
|
||||
abs_o = np.abs(omega[np.isfinite(omega)])
|
||||
vmax = float(np.percentile(abs_o, 99.5)) if abs_o.size > 0 else 1.0
|
||||
if vmax <= 0:
|
||||
vmax = 1.0
|
||||
|
||||
ny, nx = omega.shape
|
||||
fig, ax = plt.subplots(figsize=(min(18, max(8, nx / 60)), min(10, max(3, ny / 40))))
|
||||
im = ax.imshow(omega, origin="lower", aspect="equal", cmap="RdBu_r",
|
||||
vmin=-vmax, vmax=vmax, extent=(0, nx - 1, 0, ny - 1))
|
||||
ax.set_xlabel("x (lattice)")
|
||||
ax.set_ylabel("y (lattice)")
|
||||
if title:
|
||||
ax.set_title(title)
|
||||
fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label=r"$\omega_z$")
|
||||
fig.tight_layout()
|
||||
fig.savefig(path, dpi=150, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Cleanup
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def close(self):
|
||||
"""Release GPU resources."""
|
||||
self.sim.close()
|
||||
|
||||
def __del__(self):
|
||||
try:
|
||||
self.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _omega_from_nu(nu: float) -> float:
|
||||
"""Convert kinematic viscosity to relaxation parameter omega."""
|
||||
cs2 = 1.0 / 3.0
|
||||
return 1.0 / (3.0 * nu / 1.0 + 0.5)
|
||||
Reference in New Issue
Block a user