445 lines
18 KiB
Python
445 lines
18 KiB
Python
# CelerisLab/lbm/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 ..common import utils
|
|
from ..common import preprocess as preproc
|
|
from ..cuda import compiler
|
|
|
|
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,
|
|
use_kernel_v2: bool = True,
|
|
collision_model: int = 0,
|
|
streaming_model: int = 0,
|
|
store_precision: int = 0,
|
|
use_ddf_shifting: int = 0,
|
|
use_les: int = 0,
|
|
les_cs: float = 0.16,
|
|
):
|
|
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)
|
|
|
|
self.use_kernel_v2 = bool(use_kernel_v2)
|
|
self.collision_model = int(collision_model)
|
|
self.streaming_model = int(streaming_model)
|
|
self.store_precision = int(store_precision)
|
|
self.use_ddf_shifting = int(use_ddf_shifting)
|
|
self.use_les = int(use_les)
|
|
self.les_cs = float(les_cs)
|
|
|
|
if self.collision_model not in (0, 1, 2):
|
|
raise ValueError("collision_model must be 0(SRT), 1(TRT), or 2(MRT).")
|
|
if self.streaming_model not in (0, 1):
|
|
raise ValueError("streaming_model must be 0(double-buffer) or 1(esopull).")
|
|
if self.store_precision not in (0, 1, 2):
|
|
raise ValueError("store_precision must be 0(FP32), 1(FP16S), or 2(FP16C).")
|
|
if self.use_ddf_shifting not in (0, 1):
|
|
raise ValueError("use_ddf_shifting must be 0 or 1.")
|
|
if self.use_les not in (0, 1):
|
|
raise ValueError("use_les must be 0 or 1.")
|
|
if not (0.0 < self.les_cs < 1.0):
|
|
raise ValueError("les_cs must be in (0, 1).")
|
|
|
|
# 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."
|
|
)
|
|
|
|
self.objects = {}
|
|
|
|
# Compile and load kernel
|
|
self._rebuild_kernel()
|
|
|
|
# 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.action = np.zeros(0, dtype=self.DATA_TYPE)
|
|
self.obs = np.zeros(0, dtype=self.DATA_TYPE)
|
|
|
|
self.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 _configure_kernel(self):
|
|
if self.use_kernel_v2:
|
|
compiler.config_kernal_v2(
|
|
self.cuda_config,
|
|
self.field_config,
|
|
collision_model=self.collision_model,
|
|
streaming_model=self.streaming_model,
|
|
store_precision=self.store_precision,
|
|
use_ddf_shifting=self.use_ddf_shifting,
|
|
use_les=self.use_les,
|
|
les_cs=self.les_cs,
|
|
)
|
|
else:
|
|
compiler.config_kernal(self.cuda_config, self.field_config)
|
|
|
|
def _compile_and_load_kernel(self):
|
|
if self.use_kernel_v2:
|
|
compiler.compile_kernel_v2()
|
|
self.ptx = cuda.module_from_file(compiler.kernel_path("kernel_v2.ptx"))
|
|
self.step = self.ptx.get_function("OneStep")
|
|
self.initflow = self.ptx.get_function("InitTubeFlow_v2")
|
|
else:
|
|
compiler.compile_kernel()
|
|
self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx"))
|
|
self.step = self.ptx.get_function("OneStep")
|
|
self.initflow = self.ptx.get_function("InitTubeFlow")
|
|
|
|
def _rebuild_kernel(self):
|
|
self._configure_kernel()
|
|
compiler.config_object(len(self.objects))
|
|
self._compile_and_load_kernel()
|
|
|
|
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}.")
|
|
|
|
# Ensure host-side DDF mirrors current device state before local edits.
|
|
cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
|
|
|
|
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
|
|
for i in range(self.LATTICE):
|
|
self.ddf[k + i * self.FIELD_SIZE] = self.WW[i]
|
|
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)
|
|
cuda.memcpy_htod(self.ddf_gpu, self.ddf)
|
|
cuda.memcpy_htod(self.temp_gpu, self.ddf)
|
|
|
|
self._rebuild_kernel()
|
|
|
|
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,
|
|
}
|
|
|
|
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)
|
|
|
|
self._rebuild_kernel()
|
|
|
|
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 run(self, num_steps: int, action_target: np.ndarray):
|
|
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.")
|
|
elif len(self.objects) == 0:
|
|
raise ValueError("No objects have been added to the flow field.")
|
|
|
|
weight = 0.1
|
|
stream = cuda.Stream()
|
|
action_pinned = cuda.pagelocked_empty_like(self.action)
|
|
action_pinned[:] = self.action
|
|
obs_pinned = cuda.pagelocked_empty_like(self.obs)
|
|
self.obs[:] = 0
|
|
for i in range(num_steps):
|
|
action_pinned = (1 - weight) * action_pinned + weight * action_target
|
|
cuda.memcpy_htod_async(self.action_gpu, action_pinned, stream)
|
|
self.step(
|
|
self.flag_gpu,
|
|
self.ddf_gpu,
|
|
self.temp_gpu,
|
|
self.indx_gpu,
|
|
self.delta_gpu,
|
|
self.action_gpu,
|
|
self.obs_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(obs_pinned, self.obs_gpu, stream)
|
|
cuda.memset_d32_async(self.obs_gpu, 0, self.obs.size, stream)
|
|
self.obs += obs_pinned
|
|
stream.synchronize()
|
|
self.obs = (self.obs / num_steps).astype(self.DATA_TYPE)
|
|
|
|
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):
|
|
self.context.pop() |