372 lines
13 KiB
Python
372 lines
13 KiB
Python
# 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)
|