# 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 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.objects = {} self.action = np.zeros(0, dtype=self.DATA_TYPE) self.obs = np.zeros(0, dtype=self.DATA_TYPE) 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 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, } 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 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()