Frank14fandCursor
b144d62920
Replace the legacy CCD workspace with acquisition, direct-dq, original and lagged CCD contracts so the DRL-versus-constant-mean mechanism is reproducible and fail-closed. Co-authored-by: Cursor <cursoragent@cursor.com>
609 lines
28 KiB
Python
609 lines
28 KiB
Python
# CelerisLab/driver.py
|
|
|
|
import pycuda.driver as cuda
|
|
import numpy as np
|
|
import struct
|
|
from scipy.special import jv, expi
|
|
from typing import List, Tuple, Union, Optional
|
|
|
|
from . import utils
|
|
from . import preprocess as preproc
|
|
from . import compiler
|
|
from src.CCD_analysis.acquisition.solver_state import copy_ping_pong_ddf, d2q9_q_over_u0_xy
|
|
|
|
FLUID = 0b00000001
|
|
SOLID = 0b00000010
|
|
GAS = 0b00000100
|
|
INTERFACE = 0b00001000
|
|
SENSOR = 0b00010000
|
|
V_TAYLOR = np.int32(1)
|
|
|
|
class FlowField:
|
|
def __init__(
|
|
self,
|
|
field_config: utils.FlowFieldConfig,
|
|
cuda_config: utils.CudaConfig,
|
|
device_id: Union[int, List[int]] = None,
|
|
):
|
|
self.field_config = field_config
|
|
self.cuda_config = cuda_config
|
|
cuda.init()
|
|
|
|
# Sanity checks
|
|
if cuda_config.multi_gpu:
|
|
if device_id is None or isinstance(device_id, int):
|
|
raise ValueError("Multi-GPU support requires a list of device IDs.")
|
|
# self.devices = [cuda.Device(id) for id in device_id]
|
|
raise NotImplementedError("Multi-GPU support is not implemented yet.")
|
|
else:
|
|
if isinstance(device_id, list):
|
|
if len(device_id) > 1:
|
|
raise ValueError(
|
|
"Single-GPU mode does not support multiple device IDs."
|
|
)
|
|
device_id = device_id[0]
|
|
elif device_id is None:
|
|
device_id = 0
|
|
utils.check_cuda_device_availability(device_id)
|
|
self.device = cuda.Device(device_id)
|
|
self.context = self.device.make_context()
|
|
|
|
utils.check_cuda_capability(field_config, cuda_config, device_id)
|
|
|
|
# Config kernel
|
|
compiler.config_kernal(cuda_config, field_config)
|
|
compiler.config_object(int(0))
|
|
# compiler.config_sensor(int(0))
|
|
|
|
# Set constants
|
|
if field_config.data_type == "FP32":
|
|
self.DATA_TYPE = np.float32
|
|
else:
|
|
raise ValueError(f"Unsupported data type {field_config.data_type}.")
|
|
|
|
self.FIELD_SHAPE = tuple(
|
|
size * unit
|
|
for size, unit in zip(
|
|
field_config.field_dim_in_U, cuda_config.unit_dimensions
|
|
)
|
|
)
|
|
self.FIELD_SIZE = np.prod(self.FIELD_SHAPE)
|
|
self.LATTICE = field_config.lattice
|
|
self.DIM = field_config.dimensionality
|
|
if field_config.lattice == 9 and field_config.dimensionality == 2:
|
|
self.E = np.array(
|
|
[0, 0, 1, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, 1, -1, -1, 1, -1],
|
|
dtype=np.int32,
|
|
).reshape(9, 2)
|
|
self.OPP = np.array([0, 3, 4, 1, 2, 7, 8, 5, 6], dtype=np.int32)
|
|
self.WW = np.array(
|
|
[4 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 36, 1 / 36, 1 / 36, 1 / 36],
|
|
dtype=self.DATA_TYPE,
|
|
)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Unsupported lattice type {field_config.lattice} in {field_config.dimensionality} dimensions."
|
|
)
|
|
|
|
# Compile kernel
|
|
compiler.compile_kernel()
|
|
self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx"))
|
|
self.step = self.ptx.get_function("OneStep")
|
|
initflow = self.ptx.get_function("InitTubeFlow")
|
|
|
|
# Initialize memory
|
|
self.ddf = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE)
|
|
self.ddf_save = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE)
|
|
self.flag = np.ones(self.FIELD_SIZE, dtype=np.uint8)
|
|
self.indx = np.zeros(self.FIELD_SIZE, dtype=np.int32)
|
|
self.delta_curve = np.zeros(0, dtype=self.DATA_TYPE)
|
|
self.vortex_config = np.zeros(7, dtype=float)
|
|
|
|
self.ddf_gpu = cuda.mem_alloc(self.ddf.nbytes)
|
|
self.temp_gpu = cuda.mem_alloc(self.ddf.nbytes)
|
|
self.flag_gpu = cuda.mem_alloc(self.flag.nbytes)
|
|
self.indx_gpu = cuda.mem_alloc(self.indx.nbytes)
|
|
self.delta_gpu = cuda.mem_alloc(1)
|
|
self.vortex_gpu = cuda.mem_alloc(self.vortex_config.nbytes)
|
|
self.error_flag = np.zeros(1, dtype=np.uint32)
|
|
self.error_flag_gpu = cuda.mem_alloc(self.error_flag.nbytes)
|
|
self.last_error_flag = 0
|
|
|
|
self.objects = {}
|
|
self.action = np.zeros(0, dtype=self.DATA_TYPE)
|
|
self.obs = np.zeros(0, dtype=self.DATA_TYPE)
|
|
self._control_interval = None
|
|
self._last_completed_observation = None
|
|
self._last_effective_action = None
|
|
self._completed_lattice_steps = 0
|
|
self._completed_control_intervals = 0
|
|
|
|
initflow(
|
|
self.flag_gpu,
|
|
self.ddf_gpu,
|
|
block=(self.cuda_config.threads_per_block, 1, 1),
|
|
grid=(
|
|
int(self.FIELD_SHAPE[0] / self.cuda_config.threads_per_block),
|
|
int(self.FIELD_SHAPE[1]),
|
|
int(self.FIELD_SHAPE[2]),
|
|
),
|
|
)
|
|
cuda.memcpy_dtoh(self.flag, self.flag_gpu)
|
|
cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
|
|
|
|
def completed_flags_xy(self) -> np.ndarray:
|
|
"""Return a read-only copy of configured flags in canonical ``(NX, NY)`` order.
|
|
|
|
The solver stores flags flat with ``k = x + y * NX``. No bit is
|
|
interpreted or rewritten here, so FLUID/SOLID and auxiliary bits are
|
|
preserved exactly.
|
|
"""
|
|
flat = np.asarray(self.flag)
|
|
expected = int(self.FIELD_SHAPE[0]) * int(self.FIELD_SHAPE[1]) * int(self.FIELD_SHAPE[2])
|
|
if flat.dtype != np.dtype("uint8") or flat.ndim != 1 or flat.size != expected:
|
|
raise RuntimeError("configured solver flag storage is not canonical uint8 flat data")
|
|
if int(self.FIELD_SHAPE[2]) != 1:
|
|
raise RuntimeError("canonical CCD flag export supports completed D2 geometry only")
|
|
result = np.ascontiguousarray(flat.reshape((self.FIELD_SHAPE[1], self.FIELD_SHAPE[0])).T)
|
|
result.setflags(write=False)
|
|
return result
|
|
|
|
def add_cylinder(self, center: Tuple[float, float, float], radius: float, id_obj: Optional[int] = None):
|
|
x_c, y_c, z_c = center
|
|
|
|
if (
|
|
x_c - radius <= 0
|
|
or x_c + radius >= self.FIELD_SHAPE[0] - 1
|
|
or y_c - radius <= 0
|
|
or y_c + radius >= self.FIELD_SHAPE[1] - 1
|
|
):
|
|
raise ValueError("Cylinder is out of bounds.")
|
|
|
|
index = self.delta_curve.size if self.delta_curve.size > 0 else 0
|
|
|
|
if self.DATA_TYPE == np.float32:
|
|
id_object = np.int32(len(self.objects))
|
|
# max_id = max(self.objects.keys())
|
|
else:
|
|
raise ValueError(f"Unsupported data type {self.DATA_TYPE}.")
|
|
|
|
for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1):
|
|
for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1):
|
|
if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2:
|
|
k = x + y * self.FIELD_SHAPE[0]
|
|
self.flag[k] = SOLID
|
|
delta_temp = np.zeros(11, dtype=self.DATA_TYPE)
|
|
delta_temp[0] = id_object.view(self.DATA_TYPE)
|
|
for i in range(self.LATTICE):
|
|
x_neb = x + self.E[i][0]
|
|
y_neb = y + self.E[i][1]
|
|
if (x_neb - x_c) ** 2 + (y_neb - y_c) ** 2 >= radius**2:
|
|
self.flag[k] |= INTERFACE
|
|
x_i, y_i = preproc.find_circle_intersection(
|
|
x, y, x_neb, y_neb, x_c, y_c, radius
|
|
)
|
|
d_neb = np.sqrt((x_i - x_neb) ** 2 + (y_i - y_neb) ** 2)
|
|
delta_temp[i] = d_neb / np.sqrt(
|
|
self.E[i][0] ** 2 + self.E[i][1] ** 2
|
|
)
|
|
if self.flag[k] & INTERFACE:
|
|
delta_temp[9] = (y_c - y) / radius
|
|
delta_temp[10] = (x - x_c) / radius
|
|
self.delta_curve = np.concatenate(
|
|
(self.delta_curve, delta_temp)
|
|
)
|
|
self.indx[k] = index
|
|
index += delta_temp.size
|
|
|
|
self.objects[id_object] = {
|
|
"type": "cylinder",
|
|
"center": center,
|
|
"radius": radius,
|
|
}
|
|
|
|
if hasattr(self, "delta_gpu"):
|
|
self.delta_gpu.free()
|
|
self.delta_gpu = cuda.mem_alloc(self.delta_curve.nbytes)
|
|
|
|
self.action = np.zeros(len(self.objects), dtype=self.DATA_TYPE)
|
|
if hasattr(self, "action_gpu"):
|
|
self.action_gpu.free()
|
|
self.action_gpu = cuda.mem_alloc(self.action.nbytes)
|
|
|
|
self.obs = np.zeros(len(self.objects) * self.DIM, dtype=self.DATA_TYPE)
|
|
if hasattr(self, "obs_gpu"):
|
|
self.obs_gpu.free()
|
|
self.obs_gpu = cuda.mem_alloc(self.obs.nbytes)
|
|
|
|
cuda.memcpy_htod(self.delta_gpu, self.delta_curve)
|
|
cuda.memcpy_htod(self.flag_gpu, self.flag)
|
|
cuda.memcpy_htod(self.indx_gpu, self.indx)
|
|
|
|
compiler.config_object(len(self.objects))
|
|
compiler.compile_kernel()
|
|
self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx"))
|
|
self.step = self.ptx.get_function("OneStep")
|
|
|
|
def add_sensor(self, center: Tuple[float, float, float], radius: float):
|
|
x_c, y_c, z_c = center
|
|
|
|
if (
|
|
x_c - radius <= 0
|
|
or x_c + radius >= self.FIELD_SHAPE[0] - 1
|
|
or y_c - radius <= 0
|
|
or y_c + radius >= self.FIELD_SHAPE[1] - 1
|
|
):
|
|
raise ValueError("Sensor is out of bounds.")
|
|
|
|
id_object = len(self.objects)
|
|
for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1):
|
|
for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1):
|
|
if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2:
|
|
k = x + y * self.FIELD_SHAPE[0]
|
|
self.flag[k] |= SENSOR
|
|
self.indx[k] = id_object
|
|
|
|
self.objects[id_object] = {
|
|
"type": "sensor",
|
|
"center": center,
|
|
"radius": radius,
|
|
}
|
|
|
|
self.action = np.zeros(len(self.objects), dtype=self.DATA_TYPE)
|
|
if hasattr(self, "action_gpu"):
|
|
self.action_gpu.free()
|
|
self.action_gpu = cuda.mem_alloc(self.action.nbytes)
|
|
|
|
self.obs = np.zeros(len(self.objects) * self.DIM, dtype=self.DATA_TYPE)
|
|
if hasattr(self, "force_gpu"):
|
|
self.obs_gpu.free()
|
|
self.obs_gpu = cuda.mem_alloc(self.obs.nbytes)
|
|
|
|
cuda.memcpy_htod(self.flag_gpu, self.flag)
|
|
cuda.memcpy_htod(self.indx_gpu, self.indx)
|
|
|
|
compiler.config_object(len(self.objects))
|
|
compiler.compile_kernel()
|
|
self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx"))
|
|
self.step = self.ptx.get_function("OneStep")
|
|
|
|
def add_vortex(self, center: Tuple[float, float, float], radius: float, strength: float, direction: float, type: str):
|
|
x_c, y_c, z_c = center
|
|
|
|
if (
|
|
x_c - radius <= 0
|
|
or x_c + radius >= self.FIELD_SHAPE[0] - 1
|
|
or y_c - radius <= 0
|
|
or y_c + radius >= self.FIELD_SHAPE[1] - 1
|
|
):
|
|
raise ValueError("Vortex is out of bounds.")
|
|
|
|
if type not in ["lamb", "oseen", "taylor"]:
|
|
raise ValueError("Vortex type" + type + " is not supported.")
|
|
|
|
x = np.linspace(-x_c, self.FIELD_SHAPE[0] - 1 - x_c, self.FIELD_SHAPE[0])
|
|
y = np.linspace(-y_c, self.FIELD_SHAPE[1] - 1 - y_c, self.FIELD_SHAPE[1])
|
|
X, Y = np.meshgrid(x, y)
|
|
r = np.sqrt(X**2 + Y**2)
|
|
nu = self.field_config.viscosity
|
|
theta = np.arctan2(Y, X)
|
|
psi = np.zeros_like(r)
|
|
|
|
if type == "lamb":
|
|
b = 3.831705970207512
|
|
n = b / radius
|
|
u0 = strength
|
|
inside = r <= radius
|
|
outside = r > radius
|
|
|
|
psi[inside] = (2 * u0 / n / jv(0, b) * jv(1, n * r[inside]) - u0 * r[inside]) * np.sin(theta[inside])
|
|
psi[outside] = -u0 * radius**2 / r[outside] * np.sin(theta[outside])
|
|
|
|
u_vor = np.gradient(psi, axis=0)
|
|
v_vor = -np.gradient(psi, axis=1)
|
|
p_vor = -2 * (np.gradient(v_vor, axis=1) - np.gradient(u_vor, axis=0)) * psi - (u_vor**2 + v_vor**2) / 2
|
|
elif type == "oseen":
|
|
# 4 nu t = radius^2 / 4
|
|
kappa = 2 * np.pi * radius **2 * strength
|
|
u_vor = - kappa / (2 * np.pi * r) * (1 - np.exp(-4 * r**2 / radius**2)) * np.sin(theta)
|
|
v_vor = kappa / (2 * np.pi * r) * (1 - np.exp(-4 * r**2 / radius**2)) * np.cos(theta)
|
|
zeta = 4 * r**2 / radius**2
|
|
p_vor = -kappa**2 / 8 / np.pi**2 / r**2 * (-2 * zeta * (expi(-zeta) - expi(-2 * zeta)) + (1 - np.exp(-zeta))**2)
|
|
elif type == "taylor":
|
|
# 4 nu t = radius^2
|
|
M = strength * np.pi * radius**4 / 8 / nu
|
|
u_vor = - M * r * 4 * nu / radius**4 * np.exp(-r**2 / radius**2) * np.sin(theta)
|
|
v_vor = M * r * 4 * nu / radius**4 * np.exp(-r**2 / radius**2) * np.cos(theta)
|
|
p_vor = -4 * M**2 * nu**2 * np.exp(-2 * r**2 / radius**2) / np.pi**2 / radius**6
|
|
|
|
cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
|
|
ddf_temp = self.ddf.copy().reshape((self.LATTICE, self.FIELD_SHAPE[1], self.FIELD_SHAPE[0])).transpose(2, 1, 0)
|
|
u_ddf = ddf_temp[:, :, 1] + ddf_temp[:, :, 5] + ddf_temp[:, :, 8] - ddf_temp[:, :, 3] - ddf_temp[:, :, 6] - ddf_temp[:, :, 7]
|
|
v_ddf = ddf_temp[:, :, 2] + ddf_temp[:, :, 5] + ddf_temp[:, :, 6] - ddf_temp[:, :, 4] - ddf_temp[:, :, 7] - ddf_temp[:, :, 8]
|
|
p_ddf = np.sum(ddf_temp, axis=2) / 3
|
|
|
|
for i in range(self.FIELD_SHAPE[0]):
|
|
for j in range(self.FIELD_SHAPE[1]):
|
|
k = i + j * self.FIELD_SHAPE[0]
|
|
if (j == 0 or j == self.FIELD_SHAPE[1] - 1) or (i == 0 or i == self.FIELD_SHAPE[0] - 1):
|
|
continue
|
|
else:
|
|
for e in range(self.LATTICE):
|
|
u = u_ddf[i, j] + u_vor[j, i]
|
|
v = v_ddf[i, j] + v_vor[j, i]
|
|
p = p_ddf[i, j] + p_vor[j, i]
|
|
eu = self.E[e][0] * u + self.E[e][1] * v
|
|
u2 = u ** 2 + v ** 2
|
|
self.ddf[k + e * self.FIELD_SIZE] = self.WW[e] * (3 * p + 3 * eu + 4.5 * eu ** 2 - 1.5 * u2)
|
|
|
|
cuda.memcpy_htod(self.ddf_gpu, self.ddf)
|
|
|
|
# def add_vortex_gpu(self, center: Tuple[float, float, float], radius: float, strength: float, direction: float, type: str):
|
|
# x_c, y_c, z_c = center
|
|
|
|
# if (
|
|
# x_c - radius <= 0
|
|
# or x_c + radius >= self.FIELD_SHAPE[0] - 1
|
|
# or y_c - radius <= 0
|
|
# or y_c + radius >= self.FIELD_SHAPE[1] - 1
|
|
# ):
|
|
# raise ValueError("Vortex is out of bounds.")
|
|
|
|
# if type not in ["lamb", "oseen", "taylor"]:
|
|
# raise ValueError("Vortex type" + type + " is not supported.")
|
|
|
|
# add_vortex = self.ptx.get_function("AddVortex")
|
|
|
|
# self.vortex_config[0:3] = np.array(center, dtype=float)
|
|
# self.vortex_config[3] = radius
|
|
# self.vortex_config[4] = strength
|
|
# self.vortex_config[5] = direction
|
|
# if type == "taylor":
|
|
# self.vortex_config[6] =
|
|
|
|
def _validate_run(self, num_steps: int, action_target: np.ndarray):
|
|
if type(num_steps) is not int or num_steps < 1:
|
|
raise ValueError("num_steps must be a positive integer")
|
|
if action_target.size != len(self.objects) or action_target.dtype != self.DATA_TYPE:
|
|
raise ValueError("action data type or size does not match the objects.")
|
|
if len(self.objects) == 0:
|
|
raise ValueError("No objects have been added to the flow field.")
|
|
|
|
def run(self, num_steps: int, action_target: np.ndarray):
|
|
"""Advance one complete legacy interval (original public behavior)."""
|
|
if self._control_interval is not None:
|
|
raise RuntimeError("run is unavailable while a control interval is active")
|
|
self.begin_control_interval(num_steps, action_target)
|
|
self.run_control_segment(num_steps)
|
|
self.end_control_interval()
|
|
|
|
def begin_control_interval(self, total_steps: int, action_target: np.ndarray):
|
|
"""Start one policy interval that may be split only to read/save fields."""
|
|
if self._control_interval is not None:
|
|
raise RuntimeError("a control interval is already active")
|
|
self._validate_run(total_steps, action_target)
|
|
self.error_flag[0] = 0
|
|
cuda.memcpy_htod(self.error_flag_gpu, self.error_flag)
|
|
self.obs[:] = 0
|
|
action = cuda.pagelocked_empty_like(self.action)
|
|
action[:] = self.action
|
|
self._control_interval = {
|
|
"total_steps": total_steps,
|
|
"completed_steps": 0,
|
|
"target": action_target.copy(),
|
|
"action": action,
|
|
"obs_steps": cuda.pagelocked_empty((total_steps, self.obs.size), dtype=self.DATA_TYPE),
|
|
"stream": cuda.Stream(),
|
|
}
|
|
|
|
def run_control_segment(self, num_steps: int):
|
|
"""Advance part of the active interval without resetting smoothing or obs."""
|
|
state = self._control_interval
|
|
if state is None:
|
|
raise RuntimeError("no control interval is active")
|
|
if type(num_steps) is not int or num_steps < 1:
|
|
raise ValueError("num_steps must be a positive integer")
|
|
if state["completed_steps"] + num_steps > state["total_steps"]:
|
|
raise ValueError("segment exceeds the active control interval")
|
|
stream = state["stream"]
|
|
start = state["completed_steps"]
|
|
for local_step in range(num_steps):
|
|
state["action"] = 0.9 * state["action"] + 0.1 * state["target"]
|
|
cuda.memcpy_htod_async(self.action_gpu, state["action"], stream)
|
|
self.step(
|
|
self.flag_gpu,
|
|
self.ddf_gpu,
|
|
self.temp_gpu,
|
|
self.indx_gpu,
|
|
self.delta_gpu,
|
|
self.action_gpu,
|
|
self.obs_gpu,
|
|
self.error_flag_gpu,
|
|
block=(self.cuda_config.threads_per_block, 1, 1),
|
|
grid=(
|
|
int(self.FIELD_SHAPE[0] / self.cuda_config.threads_per_block),
|
|
int(self.FIELD_SHAPE[1]),
|
|
int(self.FIELD_SHAPE[2]),
|
|
),
|
|
stream=stream,
|
|
)
|
|
self.ddf_gpu, self.temp_gpu = self.temp_gpu, self.ddf_gpu
|
|
cuda.memcpy_dtoh_async(state["obs_steps"][start + local_step], self.obs_gpu, stream)
|
|
cuda.memset_d32_async(self.obs_gpu, 0, self.obs.size, stream)
|
|
stream.synchronize()
|
|
state["completed_steps"] += num_steps
|
|
self._completed_lattice_steps += num_steps
|
|
|
|
def current_step_observation(self):
|
|
"""Return raw telemetry for the latest completed lattice step.
|
|
|
|
During an active split interval this is the latest synchronized step. At
|
|
a completed control boundary it is the persisted final raw step, not the
|
|
interval-averaged public ``obs``.
|
|
"""
|
|
state = self._control_interval
|
|
if state is not None and state["completed_steps"] >= 1:
|
|
return state["obs_steps"][state["completed_steps"] - 1].copy()
|
|
if self._last_completed_observation is None:
|
|
raise RuntimeError("no completed lattice step is available")
|
|
return self._last_completed_observation.copy()
|
|
|
|
def current_effective_action(self):
|
|
"""Return the latest EMA action, including at a completed boundary."""
|
|
state = self._control_interval
|
|
if state is not None and state["completed_steps"] >= 1:
|
|
return np.asarray(state["action"]).copy()
|
|
if self._last_effective_action is None:
|
|
raise RuntimeError("no completed lattice step is available")
|
|
return self._last_effective_action.copy()
|
|
|
|
def _require_completed_split_step(self):
|
|
state = self._control_interval
|
|
if state is None or state["completed_steps"] < 1:
|
|
raise RuntimeError("no completed step is available in the active control interval")
|
|
|
|
def current_step_velocity_field(self):
|
|
"""Return completed Legacy nondimensional velocity ``q/U0`` as ``(NX, NY)``."""
|
|
if self._control_interval is not None:
|
|
self._require_completed_split_step()
|
|
elif self._last_completed_observation is None:
|
|
raise RuntimeError("no completed Legacy step is available")
|
|
# run_control_segment synchronizes before returning. After its pointer swap,
|
|
# ddf_gpu is the completed state and temp_gpu is the previous/work buffer.
|
|
cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
|
|
flags = self.completed_flags_xy()
|
|
return d2q9_q_over_u0_xy(
|
|
self.ddf, int(self.FIELD_SHAPE[0]), int(self.FIELD_SHAPE[1]), flags,
|
|
float(self.field_config.velocity),
|
|
)
|
|
|
|
def current_step_velocity_probe(self, lattice_index: Tuple[int, int]):
|
|
"""Read one synchronized Legacy nondimensional ``(ux/U0, uy/U0)`` pair."""
|
|
self._require_completed_split_step()
|
|
if (not isinstance(lattice_index, tuple) or len(lattice_index) != 2
|
|
or any(type(value) is not int for value in lattice_index)):
|
|
raise ValueError("lattice_index must be an (x, y) integer tuple")
|
|
x, y = lattice_index
|
|
if not (0 <= x < self.FIELD_SHAPE[0] and 0 <= y < self.FIELD_SHAPE[1]):
|
|
raise ValueError("velocity probe is outside the lattice")
|
|
ux, uy = self.current_step_velocity_field()
|
|
return np.asarray([ux[x, y], uy[x, y]], dtype=self.DATA_TYPE)
|
|
|
|
def active_step_clock_state(self):
|
|
"""Return solver lineage during a split interval after a completed step.
|
|
|
|
The control clock is the number of fully completed control intervals; it
|
|
therefore identifies the active interval's zero-based absolute index.
|
|
"""
|
|
state = self._control_interval
|
|
if state is None or state["completed_steps"] < 1:
|
|
raise RuntimeError("active-step clocks require a split interval with a completed step")
|
|
return {
|
|
"solver_absolute_lattice_clock": int(self._completed_lattice_steps),
|
|
"solver_absolute_control_clock": int(self._completed_control_intervals),
|
|
}
|
|
|
|
def solver_clock_state(self):
|
|
"""Return public absolute solver lineage clocks at the current boundary."""
|
|
if self._control_interval is not None:
|
|
raise RuntimeError("solver clocks are boundary-safe only")
|
|
return {
|
|
"solver_absolute_lattice_clock": int(self._completed_lattice_steps),
|
|
"solver_absolute_control_clock": int(self._completed_control_intervals),
|
|
}
|
|
|
|
def full_state_checkpoint(self):
|
|
"""Capture exact restart state only at a completed control boundary."""
|
|
if self._control_interval is not None:
|
|
raise RuntimeError("full checkpoint requires a completed control boundary")
|
|
ddf = self.current_step_ddf_checkpoint()
|
|
return {
|
|
**ddf,
|
|
"action": self.action.copy(),
|
|
"last_effective_action": None if self._last_effective_action is None else self._last_effective_action.copy(),
|
|
"raw_observation": None if self._last_completed_observation is None else self._last_completed_observation.copy(),
|
|
"boundary_observation": self.obs.copy(),
|
|
"solver_absolute_lattice_clock": int(self._completed_lattice_steps),
|
|
"solver_absolute_control_clock": int(self._completed_control_intervals),
|
|
}
|
|
|
|
def restore_full_state(self, checkpoint):
|
|
"""Restore both ping-pong DDFs and solver-side boundary lifecycle state."""
|
|
if self._control_interval is not None:
|
|
raise RuntimeError("cannot restore during an active control interval")
|
|
current = np.asarray(checkpoint["current_ddf"]); temp = np.asarray(checkpoint["temp_ddf"])
|
|
action = np.asarray(checkpoint["action"]); boundary = np.asarray(checkpoint["boundary_observation"])
|
|
raw = checkpoint["raw_observation"]; effective = checkpoint["last_effective_action"]
|
|
if current.dtype != self.DATA_TYPE or temp.dtype != self.DATA_TYPE or current.shape != self.ddf.shape or temp.shape != self.ddf.shape:
|
|
raise ValueError("checkpoint DDF storage mismatch")
|
|
if action.dtype != self.DATA_TYPE or action.shape != self.action.shape or boundary.dtype != self.DATA_TYPE or boundary.shape != self.obs.shape:
|
|
raise ValueError("checkpoint action/observation mismatch")
|
|
if raw is not None and (np.asarray(raw).dtype != self.DATA_TYPE or np.asarray(raw).shape != self.obs.shape): raise ValueError("checkpoint raw observation mismatch")
|
|
if effective is not None and (np.asarray(effective).dtype != self.DATA_TYPE or np.asarray(effective).shape != self.action.shape): raise ValueError("checkpoint effective action mismatch")
|
|
cuda.memcpy_htod(self.ddf_gpu, current); cuda.memcpy_htod(self.temp_gpu, temp)
|
|
self.ddf = current.copy(); self.action = action.copy(); self.obs = boundary.copy()
|
|
self._last_completed_observation = None if raw is None else np.asarray(raw).copy()
|
|
self._last_effective_action = None if effective is None else np.asarray(effective).copy()
|
|
self._completed_lattice_steps = int(checkpoint["solver_absolute_lattice_clock"])
|
|
self._completed_control_intervals = int(checkpoint["solver_absolute_control_clock"])
|
|
cuda.memcpy_htod(self.action_gpu, self.action)
|
|
|
|
def current_step_ddf_checkpoint(self):
|
|
"""Return copies/hashes of current(completed) and temp(previous/work) buffers.
|
|
|
|
Synchronous device-to-host copies make this safe both at a completed split
|
|
step and at a completed control boundary; no solver state is modified.
|
|
"""
|
|
return copy_ping_pong_ddf(
|
|
lambda host: cuda.memcpy_dtoh(host, self.ddf_gpu),
|
|
lambda host: cuda.memcpy_dtoh(host, self.temp_gpu),
|
|
int(self.FIELD_SIZE * self.LATTICE),
|
|
)
|
|
|
|
def end_control_interval(self):
|
|
"""Publish obs once, only at the original policy-control boundary."""
|
|
state = self._control_interval
|
|
if state is None:
|
|
raise RuntimeError("no control interval is active")
|
|
if state["completed_steps"] != state["total_steps"]:
|
|
raise RuntimeError("cannot end a control interval before its boundary")
|
|
self.obs[:] = 0
|
|
for step_obs in state["obs_steps"]:
|
|
self.obs += step_obs
|
|
self.obs = (self.obs / state["total_steps"]).astype(self.DATA_TYPE)
|
|
cuda.memcpy_dtoh(self.error_flag, self.error_flag_gpu)
|
|
self.last_error_flag = int(self.error_flag[0])
|
|
# Persist the final lattice-step state before clearing the split lifecycle.
|
|
# The next begin_control_interval therefore starts its EMA from this exact
|
|
# action, preserving the historical uninterrupted-run semantics.
|
|
self.action = np.asarray(state["action"], dtype=self.DATA_TYPE).copy()
|
|
self._last_effective_action = self.action.copy()
|
|
self._last_completed_observation = state["obs_steps"][-1].copy()
|
|
self._completed_control_intervals += 1
|
|
self._control_interval = None
|
|
return self.obs
|
|
|
|
def has_numeric_error(self) -> bool:
|
|
return bool(self.last_error_flag != 0)
|
|
|
|
def apply_ddf(self):
|
|
cuda.memcpy_htod(self.ddf_gpu, self.ddf)
|
|
|
|
def get_ddf(self):
|
|
cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
|
|
|
|
def save_ddf(self):
|
|
self.ddf_save = self.ddf.copy()
|
|
|
|
def restore_ddf(self):
|
|
self.ddf = self.ddf_save.copy()
|
|
|
|
def __del__(self):
|
|
# Shutdown order can invalidate current context before object cleanup.
|
|
ctx = getattr(self, "context", None)
|
|
if ctx is None:
|
|
return
|
|
try:
|
|
ctx.pop()
|
|
except Exception:
|
|
pass |