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