Major update: functional LBM and DRL

.gitignore

@@ -0,0 +1,2 @@
+tensorboard/*
+models/*

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "[cuda-cpp]": {
+    },
+    "C_Cpp.errorSquiggles": "disabled"
+}

CUDA/kernel.cu

@@ -1,86 +0,0 @@
-#include "setting.h"
-
-__global__ void Collision(int *flag, LBtype *pres, LBtype *vell, LBtype *f0, LBtype *forc)
-{
-	GLOBAL_INDEX()
-
-	if(flag[k]==0)
-	{
-		LBtype P,Ux,Uy;
-		LBtype M[9];
-		LBtype g[9];
-
-		LBtype Fx=forc[k*2], Fy=forc[k*2+1];
-
-		for(int kk=0;kk<9;kk++)
-			g[kk]=f0[k*9+kk];
-
-		Ux=(g[1]+g[5]+g[8]-g[3]-g[6]-g[7]+0.5*Fx)/rho;
-		Uy=(g[2]+g[5]+g[6]-g[4]-g[7]-g[8]+0.5*Fy)/rho;
-		P =(g[0]+g[1]+g[2]+g[3]+g[4]+g[5]+g[6]+g[7]+g[8])/3.0;
-		pressure[k]=P;
-
-		M[0]=   g[0]  +g[1]  +g[2]  +g[3]  +g[4]  +g[5]  +g[6]  +g[7]  +g[8];
-		M[1]=-4*g[0]  -g[1]  -g[2]  -g[3]  -g[4]+2*g[5]+2*g[6]+2*g[7]+2*g[8];
-		M[2]= 4*g[0]-2*g[1]-2*g[2]-2*g[3]-2*g[4]  +g[5]  +g[6]  +g[7]  +g[8];
-		M[3]=          g[1]         -g[3]         +g[5]  -g[6]  -g[7]  +g[8];
-		M[4]=       -2*g[1]       +2*g[3]         +g[5]  -g[6]  -g[7]  +g[8];
-		M[5]=                 g[2]         -g[4]  +g[5]  +g[6]  -g[7]  -g[8];
-		M[6]=              -2*g[2]       +2*g[4]  +g[5]  +g[6]  -g[7]  -g[8];
-		M[7]=          g[1]  -g[2]  +g[3]  -g[4];
-		M[8]=                                      g[5]  -g[6]  +g[7]  -g[8];	
-
-		M[0]=1.00*( 3*P	                   -M[0]);
-		M[1]=1.20*(-6*P+3*rho*(Ux*Ux+Uy*Uy)-M[1])+(1-0.5*1.2)*6*(Ux*Fx+Uy*Fy);
-		M[2]=1.20*( 3*P-3*rho*(Ux*Ux+Uy*Uy)-M[2])-(1-0.5*1.2)*6*(Ux*Fx+Uy*Fy);
-		M[3]=1.00*( rho*Ux                 -M[3])+(1-0.5*0.0)*Fx;
-		M[4]=1.20*(-rho*Ux                 -M[4])-(1-0.5*1.2)*Fx;
-		M[5]=1.00*( rho*Uy                 -M[5])+(1-0.5*0.0)*Fy;
-		M[6]=1.20*(-rho*Uy                 -M[6])-(1-0.5*1.2)*Fy;
-		M[7]=  nu*(rho*(Ux*Ux-Uy*Uy)       -M[7])+(1-0.5*nu)*2*(Ux*Fx-Uy*Fy);
-		M[8]=  nu*(rho*Ux*Uy               -M[8])+(1-0.5*nu)*(Ux*Fy+Uy*Fx);
-
-		f0[k*9]  =g[0]+(  M[0]  -M[1]  +M[2])/9.0;
-		f0[k*9+1]=g[1]+(4*M[0]  -M[1]-2*M[2]+6*M[3]-6*M[4]              +9*M[7])/36.0;
-		f0[k*9+2]=g[2]+(4*M[0]  -M[1]-2*M[2]              +6*M[5]-6*M[6]-9*M[7])/36.0;
-		f0[k*9+3]=g[3]+(4*M[0]  -M[1]-2*M[2]-6*M[3]+6*M[4]              +9*M[7])/36.0;
-		f0[k*9+4]=g[4]+(4*M[0]  -M[1]-2*M[2]              -6*M[5]+6*M[6]-9*M[7])/36.0;
-		f0[k*9+5]=g[5]+(4*M[0]+2*M[1]  +M[2]+6*M[3]+3*M[4]+6*M[5]+3*M[6]       +9*M[8])/36.0;
-		f0[k*9+6]=g[6]+(4*M[0]+2*M[1]  +M[2]-6*M[3]-3*M[4]+6*M[5]+3*M[6]       -9*M[8])/36.0;
-		f0[k*9+7]=g[7]+(4*M[0]+2*M[1]  +M[2]-6*M[3]-3*M[4]-6*M[5]-3*M[6]       +9*M[8])/36.0;
-		f0[k*9+8]=g[8]+(4*M[0]+2*M[1]  +M[2]+6*M[3]+3*M[4]-6*M[5]-3*M[6]       -9*M[8])/36.0;
-	}
-}
-
-__global__ void Streaming(LBtype *f0, LBtype *f1)
-{
-	GLOBAL_INDEX()
-	int neighbor,nex,ney;
-	int e[9][2]={{0,0},{1,0},{0,1},{-1,0},{0,-1},{1,1},{-1,1},{-1,-1},{1,-1}};
-
-	for(int kk=0;kk<9;kk++)
-	{
-		nex=(x+e[kk][0]+NX)%NX;
-		ney=(y+e[kk][1]+NY)%NY;
-		neighbor=ney*NX+nex;
-		f1[neighbor*9+kk]=f0[k*9+kk];
-	}
-}
-
-__global__ void BounceBack(int *flag, LBtype *f0)
-{
-	GLOBAL_INDEX()
-	int neighbor,nex,ney;
-	int e[9][2]={{0,0},{1,0},{0,1},{-1,0},{0,-1},{1,1},{-1,1},{-1,-1},{1,-1}};
-	int opp[9]={0,3,4,1,2,7,8,5,6};
-
-	if(flag[k]==1)
-		for(int kk=1;kk<9;kk++)
-		{
-			nex=(x+e[kk][0]+NX)%NX;
-			ney=(y+e[kk][1]+NY)%NY;
-			neighbor=ney*NX+nex;
-			if(flag[neighbor]==0)
-				f0[neighbor*9+kk]=f0[k*9+opp[kk]];
-		}
-}

CUDA/setting.h

@@ -1,9 +0,0 @@
-#define LBtype double
-#define Pi 3.141592653589793238
-const int N_thread=256;
-int devicenum=0;
-
-#define GLOBAL_INDEX() \
-    int x = threadIdx.x + blockDim.x * blockIdx.x; \
-    int y = blockIdx.y; \
-    int k = y * gridDim.x * blockDim.x + x;

CelerisLab/__init__.py

@@ -0,0 +1,3 @@
+# CelerisLab/__init__.py
+
+from .driver import FlowField

CelerisLab/compiler.py

@@ -0,0 +1,87 @@
+# CelerisLab/kernels/compiler.py
+
+import subprocess
+import re
+import os
+
+from .utils import FlowFieldConfig, CudaConfig
+
+
+def kernel_path(file_name: str) -> str:
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    return os.path.join(current_dir, "kernels", file_name)
+
+
+def read_lines(file_path):
+    with open(file_path, "r") as file:
+        lines = file.readlines()
+    return lines
+
+
+def write_lines(file_path, lines):
+    with open(file_path, "w") as file:
+        file.writelines(lines)
+
+
+def modify_macro(lines, macro_name, new_value):
+    pattern = re.compile(rf"(#define\s+{macro_name}\s+)(\S+)")
+    for i, line in enumerate(lines):
+        if pattern.match(line):
+            lines[i] = pattern.sub(rf"\g<1>{new_value}", line)
+            break
+    return lines
+
+def modify_const(lines, const_name, new_type, new_value):
+    pattern = re.compile(rf"(__constant__\s+)(\S+\s+{const_name}\s*=\s*)([^;]+)(;)")
+    for i, line in enumerate(lines):
+        if pattern.match(line):
+            lines[i] = pattern.sub(rf"\g<1>{new_type} {const_name} = {new_value}\4", line)
+            break
+    return lines
+
+def compile_kernel():
+    subprocess.run(
+        [
+            "nvcc",
+            "-ptx",
+            kernel_path("kernel.cu"),
+            "-o",
+            kernel_path("kernel.ptx"),
+        ]
+    )
+
+def config_kernal(config_cuda: CudaConfig, config_field: FlowFieldConfig):
+    lines = read_lines(kernel_path("macros.h"))
+    lines = modify_macro(lines, "MULT_GPU", config_cuda.multi_gpu)
+    lines = modify_macro(lines, "NT", config_cuda.threads_per_block)
+    lines = modify_macro(lines, "X_1U", config_cuda.unit_dimensions[0])
+    lines = modify_macro(lines, "Y_1U", config_cuda.unit_dimensions[1])
+    lines = modify_macro(lines, "Z_1U", config_cuda.unit_dimensions[2])
+
+    if config_field.data_type == "FP32":
+        lb_type = "float"
+    else:
+        raise ValueError(f"Unsupported data type {config_field.data_type}.")
+    lines = modify_macro(lines, "LBtype", lb_type)
+    lines = modify_macro(lines, "UX", config_field.field_dim_in_U[0])
+    lines = modify_macro(lines, "UY", config_field.field_dim_in_U[1])
+    lines = modify_macro(lines, "UZ", config_field.field_dim_in_U[2])
+    lines = modify_macro(lines, "NX", config_field.field_dim_in_U[0] * config_cuda.unit_dimensions[0])
+    lines = modify_macro(lines, "NY", config_field.field_dim_in_U[1] * config_cuda.unit_dimensions[1])
+    lines = modify_macro(lines, "NZ", config_field.field_dim_in_U[2] * config_cuda.unit_dimensions[2])
+    lines = modify_macro(lines, "DIM", config_field.dimensionality)
+    lines = modify_macro(lines, "NQ", config_field.lattice)
+    lines = modify_macro(lines, "VIS", config_field.viscosity)
+    lines = modify_macro(lines, "U0", config_field.velocity)
+
+    write_lines(kernel_path("macros.h"), lines)
+
+def config_object(n_obj: int):
+    lines = read_lines(kernel_path("macros.h"))
+    lines = modify_macro(lines, "N_OBJS", n_obj)
+    write_lines(kernel_path("macros.h"), lines)
+
+def config_sensor(n_sen: int):
+    lines = read_lines(kernel_path("macros.h"))
+    lines = modify_macro(lines, "N_SENS", n_sen)
+    write_lines(kernel_path("macros.h"), lines)

CelerisLab/driver.py

@@ -0,0 +1,282 @@
+# CelerisLab/driver.py
+
+import pycuda.driver as cuda
+import numpy as np
+from typing import List, Tuple, Union
+
+from . import utils
+from . import preprocess as preproc
+from . import compiler
+
+FLUID = 0b00000001
+SOLID = 0b00000010
+GAS = 0b00000100
+INTERFACE = 0b00001000
+SENSOR = 0b00010000
+
+
+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.velo = np.zeros(self.FIELD_SIZE * self.DIM, 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.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.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):
+        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))
+        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, "force_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 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.")
+
+        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 __del__(self):
+        self.context.pop()

CelerisLab/kernels/D2Q9.cu

@@ -0,0 +1,101 @@
+#include "macros.h"
+#include "const.h"
+
+__device__ void Index_lattice(int &x, int &y, int &k) {
+	// Only for D2
+	x = threadIdx.x + NT * blockIdx.x;
+	y = blockIdx.y;
+	k = y * NX + x;
+}
+
+__device__ void CollisionKernel(LBtype* g, LBtype* m) {
+	// Only for D2Q9
+	LBtype p, u, v;
+    LBtype niu = 1.0 / (0.5 + 3 * VIS);
+
+	u = (g[1]+g[5]+g[8]-g[3]-g[6]-g[7])/RHO;
+	v = (g[2]+g[5]+g[6]-g[4]-g[7]-g[8])/RHO;
+	p = (g[0]+g[1]+g[2]+g[3]+g[4]+g[5]+g[6]+g[7]+g[8])/3.0;
+
+	m[0]=   g[0]  +g[1]  +g[2]  +g[3]  +g[4]  +g[5]  +g[6]  +g[7]  +g[8];
+	m[1]=-4*g[0]  -g[1]  -g[2]  -g[3]  -g[4]+2*g[5]+2*g[6]+2*g[7]+2*g[8];
+	m[2]= 4*g[0]-2*g[1]-2*g[2]-2*g[3]-2*g[4]  +g[5]  +g[6]  +g[7]  +g[8];
+	m[3]=          g[1]         -g[3]         +g[5]  -g[6]  -g[7]  +g[8];
+	m[4]=       -2*g[1]       +2*g[3]         +g[5]  -g[6]  -g[7]  +g[8];
+	m[5]=                 g[2]         -g[4]  +g[5]  +g[6]  -g[7]  -g[8];
+	m[6]=              -2*g[2]       +2*g[4]  +g[5]  +g[6]  -g[7]  -g[8];
+	m[7]=          g[1]  -g[2]  +g[3]  -g[4];
+	m[8]=                                      g[5]  -g[6]  +g[7]  -g[8];	
+	
+	m[0]=1.00*( 3*p 	            -m[0]);
+	m[1]=1.20*(-6*p +3*RHO*(u*u+v*v)-m[1]);
+	m[2]=1.20*( 3*p -3*RHO*(u*u+v*v)-m[2]);
+	m[3]=1.00*( RHO*u               -m[3]);
+	m[4]=1.20*(-RHO*u               -m[4]);
+	m[5]=1.00*( RHO*v               -m[5]);
+	m[6]=1.20*(-RHO*v               -m[6]);
+	m[7]= niu*( RHO*(u*u-v*v)       -m[7]);
+	m[8]= niu*( RHO*u*v             -m[8]);
+
+	g[0]=g[0]+(  m[0]  -m[1]  +m[2]                                   )/ 9.0;
+	g[1]=g[1]+(4*m[0]  -m[1]-2*m[2]+6*m[3]-6*m[4]              +9*m[7])/36.0;
+	g[2]=g[2]+(4*m[0]  -m[1]-2*m[2]              +6*m[5]-6*m[6]-9*m[7])/36.0;
+	g[3]=g[3]+(4*m[0]  -m[1]-2*m[2]-6*m[3]+6*m[4]              +9*m[7])/36.0;
+	g[4]=g[4]+(4*m[0]  -m[1]-2*m[2]              -6*m[5]+6*m[6]-9*m[7])/36.0;
+	g[5]=g[5]+(4*m[0]+2*m[1]  +m[2]+6*m[3]+3*m[4]+6*m[5]+3*m[6]+9*m[8])/36.0;
+	g[6]=g[6]+(4*m[0]+2*m[1]  +m[2]-6*m[3]-3*m[4]+6*m[5]+3*m[6]-9*m[8])/36.0;
+	g[7]=g[7]+(4*m[0]+2*m[1]  +m[2]-6*m[3]-3*m[4]-6*m[5]-3*m[6]+9*m[8])/36.0;
+	g[8]=g[8]+(4*m[0]+2*m[1]  +m[2]+6*m[3]+3*m[4]-6*m[5]-3*m[6]-9*m[8])/36.0;
+}
+
+__device__ void ParabolicInlet(LBtype* f, LBtype* f_neb, LBtype y) {
+	LBtype p, u, v, yy;
+	LBtype feq1, feq5, feq8, feqn1, feqn5, feqn8;
+
+	p=(f_neb[0]+f_neb[1]+f_neb[2]+f_neb[3]+f_neb[4]+f_neb[5]+f_neb[6]+f_neb[7]+f_neb[8])/3.0;
+	yy=(y-0.5*(NY-1))/(NY-2.0);
+	u=U0*1.5*(1-4*yy*yy);
+	v=0.0;
+
+	feq1=(2*p+RHO*(2*u*u+2*u    -v*v)  )/ 6.0;
+	feq5=(  p+RHO*(  u*u+3*u*v+u+v*v+v))/12.0;
+	feq8=(  p+RHO*(  u*u-3*u*v+u+v*v-v))/12.0;
+
+	u=(f_neb[1]+f_neb[5]+f_neb[8]-f_neb[3]-f_neb[6]-f_neb[7])/RHO;
+	v=(f_neb[2]+f_neb[5]+f_neb[6]-f_neb[4]-f_neb[7]-f_neb[8])/RHO;
+
+	feqn1=(2*p+RHO*(2*u*u+2*u    -v*v)  )/ 6.0;
+	feqn5=(  p+RHO*(  u*u+3*u*v+u+v*v+v))/12.0;
+	feqn8=(  p+RHO*(  u*u-3*u*v+u+v*v-v))/12.0;
+
+	f[1]=f_neb[1]-feqn1+feq1;
+	f[5]=f_neb[5]-feqn5+feq5;
+	f[8]=f_neb[8]-feqn8+feq8;
+}
+
+__device__ void PressureOutlet(LBtype* f, LBtype* f_neb, LBtype y) {
+	// Edit to Parabolic Outlet temporarily
+	LBtype p, u, v, yy;	
+	LBtype feq3, feq6, feq7, feqn3, feqn6, feqn7;
+
+	p=0.0;
+
+	yy=(y-0.5*(NY-1))/(NY-2.0);
+	u=U0*1.5*(1-4*yy*yy);
+	v=0.0;
+
+	feq3=(2*p-RHO*(-2*u*u+2*u    +v*v)  )/ 6.0;
+	feq6=(  p+RHO*(   u*u-3*u*v-u+v*v+v))/12.0;
+	feq7=(  p+RHO*(   u*u+3*u*v-u+v*v-v))/12.0;
+
+	u=(f_neb[1]+f_neb[5]+f_neb[8]-f_neb[3]-f_neb[6]-f_neb[7])/RHO;
+	v=(f_neb[2]+f_neb[5]+f_neb[6]-f_neb[4]-f_neb[7]-f_neb[8])/RHO;
+	// p=(f_neb[0]+f_neb[1]+f_neb[2]+f_neb[3]+f_neb[4]+f_neb[5]+f_neb[6]+f_neb[7]+f_neb[8])/3.0;
+	feqn3=(2*p-RHO*(-2*u*u+2*u    +v*v)  )/ 6.0;
+	feqn6=(  p+RHO*(   u*u-3*u*v-u+v*v+v))/12.0;
+	feqn7=(  p+RHO*(   u*u+3*u*v-u+v*v-v))/12.0;
+
+	f[3]=f_neb[3]-feqn3+feq3;
+	f[6]=f_neb[6]-feqn6+feq6;
+	f[7]=f_neb[7]-feqn7+feq7;
+}

CelerisLab/kernels/const.h

@@ -0,0 +1,10 @@
+// CelerisLab/kernels/const.h
+
+#ifndef CONST_H
+#define CONST_H
+
+__constant__ int e[9][2] = {{0, 0}, {1, 0}, {0, 1}, {-1, 0}, {0, -1}, {1, 1}, {-1, 1}, {-1, -1}, {1, -1}};
+__constant__ int opp[9] = {0, 3, 4, 1, 2, 7, 8, 5, 6};
+__constant__ float w[9] = {4/9., 1/9., 1/9., 1/9., 1/9., 1/36., 1/36., 1/36., 1/36.};
+
+#endif

CelerisLab/kernels/kernel.cu

@@ -0,0 +1,190 @@
+// CelerisLab/kernels/kernel.cu
+
+#include <stdio.h>
+#include <stdint.h>
+#include <cuda.h>
+
+#include "macros.h"
+#include "const.h"
+#include "D2Q9.cu"
+
+extern "C"
+{
+	__global__ void OneStep(uint8_t *flag, LBtype *f, LBtype *f_temp, int32_t *indx, LBtype *delta, LBtype *action, LBtype *obs)
+	{
+		__shared__ LBtype f_share[NT * NQ];
+		__shared__ LBtype obs_share[(N_OBJS * DIM > 0) ? N_OBJS * DIM : 1];
+
+		int x, y, k;
+		LBtype g[NQ], m[NQ];
+		Index_lattice(x, y, k); // Only for D2
+		int totalCells = NX * NY;
+		int id = indx[k];
+
+		for (int i = 0; i < NQ; i++)
+		{
+			f_share[threadIdx.x + i * NT] = f[k + i * totalCells];
+		}
+		for (int i = threadIdx.x; i < N_OBJS * DIM; i+=NT)
+		{
+			obs_share[i] = 0;
+		}
+
+		__syncthreads();
+
+		for (int i = 0; i < NQ; i++)
+		{
+			g[i] = f_share[threadIdx.x + i * NT];
+		}
+
+		if (flag[k] & FLUID)
+		{
+			CollisionKernel(g, m);
+
+			for (int i = 0; i < NQ; i++)
+			{
+				f_share[threadIdx.x + i * NT] = g[i];
+			}
+		}
+		else if (flag[k] & SOLID)
+		{
+			if (x == 0)
+			{
+				for (int i = 0; i < NQ; i++)
+				{
+					m[i] = f_share[threadIdx.x + i * NT + 1];
+				}
+				ParabolicInlet(g, m, y);
+			}
+			else if (x == NX - 1)
+			{
+				for (int i = 0; i < NQ; i++)
+				{
+					m[i] = f_share[threadIdx.x + i * NT - 1];
+				}
+				PressureOutlet(g, m, y);
+			}
+
+			for (int i = 0; i < NQ; i++)
+			{
+				f_share[threadIdx.x + i * NT] = g[i];
+			}
+		}
+
+		__syncthreads();
+
+		for (int i = 0; i < NQ; i++)
+		{
+			int x_neb = x + e[i][0];
+			int y_neb = y + e[i][1];
+
+			if (y != 0 && y != NY - 1)
+			{
+				if ((y == 1 && y_neb == 0) || (y == NY - 2 && y_neb == NY - 1))
+				{
+					f_temp[k + opp[i] * totalCells] = f_share[threadIdx.x + i * NT];
+				}
+				else
+				{
+					int k_neb = ((y_neb * NX + x_neb) + totalCells) % totalCells;
+					f_temp[k_neb + i * totalCells] = f_share[threadIdx.x + i * NT];
+				}
+			}
+		}
+
+		__syncthreads();
+
+		if (flag[k] & SOLID && flag[k] & INTERFACE)
+		{
+			LBtype Uw, Vw;
+			int id_obj = *reinterpret_cast<int*>(&delta[id]);
+			Uw = action[id_obj] * delta[id + 9];
+			Vw = action[id_obj] * delta[id + 10];
+
+			int x_neb, y_neb, k_neb;
+			for (int i = 1; i < 9; i++)
+			{
+				x_neb = x + e[i][0];
+				y_neb = y + e[i][1];
+				k_neb = x_neb + y_neb * NX;
+				if (flag[k_neb] & FLUID)
+				{
+					LBtype q = delta[id + i];
+					int k_neb2 = (y + 2 * e[i][1]) * NX + (x + 2 * e[i][0]);
+					LBtype temp = 6 * w[i] * (e[i][0] * Uw + e[i][1] * Vw);
+					f_temp[k_neb + i * totalCells] = (q * f_temp[k + opp[i] * totalCells] \
+						+ (1 - q) * f_temp[k_neb + opp[i] * totalCells] \
+						+ q * f_temp[k_neb2 + i * totalCells] + temp) / (1 + q);
+					f_temp[k + i * totalCells] = temp * Uw;
+					k_neb2 = (y - e[i][1]) * NX + (x - e[i][0]);
+					f_temp[k_neb2 + i * totalCells] = temp * Vw;
+
+					temp = f_temp[k_neb + i * totalCells] + f_temp[k + opp[i] * totalCells];
+					k_neb2 = (y - e[i][1]) * NX + (x - e[i][0]);
+					atomicAdd(&obs_share[DIM * id_obj], -temp * e[i][0] + f_temp[k + i * totalCells]);
+					atomicAdd(&obs_share[DIM * id_obj + 1], -temp * e[i][1] + f_temp[k_neb2 + i * totalCells]);
+				}
+			}
+		}
+		if (flag[k] & SENSOR)
+		{
+			LBtype u, v;
+			u = (g[1]+g[5]+g[8]-g[3]-g[6]-g[7])/RHO;
+			v = (g[2]+g[5]+g[6]-g[4]-g[7]-g[8])/RHO;
+			atomicAdd(&obs_share[DIM * id], u);
+			atomicAdd(&obs_share[DIM * id + 1], v);
+		}
+
+		__syncthreads();
+
+		for (int i = threadIdx.x; i < N_OBJS * DIM; i+=NT)
+		{
+			atomicAdd(&obs[i], obs_share[i]);
+		}
+	}
+
+	__global__ void InitTubeFlow(uint8_t *flag, LBtype *f)
+	{
+		__shared__ LBtype f_share[NT * NQ];
+		__shared__ uint8_t flag_share[NT];
+		int x, y, k;
+		LBtype u;
+		Index_lattice(x, y, k);
+		int totalCells = NX * NY;
+
+		flag_share[threadIdx.x] = flag[k];
+		for (int i = 0; i < NQ; i++)
+		{
+			f_share[threadIdx.x + i * NT] = f[k + i * totalCells];
+		}
+
+		__syncthreads();
+
+		u = U0 * 1.5 * (1 - 4 * (y - 0.5 * (NY - 1)) * (y - 0.5 * (NY - 1)) / ((NY - 2) * (NY - 2)));
+		if (y == 0 || y == NY - 1 || x == 0 || x == NX - 1)
+		{
+			flag_share[threadIdx.x] = SOLID;
+			for (int i = 0; i < NQ; i++)
+			{
+				f_share[threadIdx.x + i * NT] = 0;
+			}
+		}
+		else
+		{
+			flag_share[threadIdx.x] = FLUID;
+			for (int i = 0; i < NQ; i++)
+			{
+				f_share[threadIdx.x + i * NT] = w[i] * RHO * (3 * e[i][0] * u + \
+					4.5 * e[i][0] * e[i][0] * u * u - 1.5 * u * u);
+			}
+		}
+
+		__syncthreads();
+
+		flag[k] = flag_share[threadIdx.x];
+		for (int i = 0; i < NQ; i++)
+		{
+			f[k + i * totalCells] = f_share[threadIdx.x + i * NT];
+		}
+	}
+}

CelerisLab/kernels/macros.h

@@ -0,0 +1,34 @@
+// CelerisLab/kernels/macros.h
+
+// cuda parameters
+#define MULT_GPU False
+#define NT 128
+#define X_1U 128
+#define Y_1U 32
+#define Z_1U 1
+
+// flow parameters
+#define LBtype float
+#define UX 10
+#define UY 16
+#define UZ 1
+#define NX 1280
+#define NY 512
+#define NZ 1
+#define DIM 2
+#define NQ 9
+#define VIS 0.004
+#define RHO 1.0
+#define U0 0.01
+
+// constants
+#define PI 3.141592653589793238
+#define FLUID       0b00000001
+#define SOLID       0b00000010
+#define GAS         0b00000100
+#define INTERFACE   0b00001000
+#define SENSOR      0b00010000
+
+// variables
+#define N_OBJS 7
+// #define N_SENS 2

CelerisLab/kernels/preproc.cu

@@ -0,0 +1,2 @@
+#include "macros.h"
+#include "const.h"

CelerisLab/preprocess.py

@@ -0,0 +1,40 @@
+# CelerisLab/preprocess.py
+
+import math
+import numpy as np
+from typing import Tuple
+
+FLUID = 0b00000001
+SOLID = 0b00000010
+GAS = 0b00000100
+INTERFACE = 0b00001000
+SENSOR = 0b00010000
+
+
+def find_circle_intersection(x, y, x_neb, y_neb, xc, yc, r0):
+    dx, dy = x_neb - x, y_neb - y
+    a = dx ** 2 + dy ** 2
+    b = 2 * (dx * (x - xc) + dy * (y - yc))
+    c = (x - xc) ** 2 + (y - yc) ** 2 - r0 ** 2
+    det = b ** 2 - 4 * a * c
+
+    if det < 0:
+        return None
+
+    t1 = (-b + math.sqrt(det)) / (2 * a)
+    t2 = (-b - math.sqrt(det)) / (2 * a)
+
+    if 0 <= t1 <= 1:
+        return x + t1 * dx, y + t1 * dy
+    elif 0 <= t2 <= 1:
+        return x + t2 * dx, y + t2 * dy
+    else:
+        return None 
+    
+def find_sensor_area(radius):
+    area = 0
+    for i in range(np.floor(-radius), np.ceil(radius)):
+        for j in range(np.floor(-radius), np.ceil(radius)):
+            if i ** 2 + j ** 2 <= radius ** 2:
+                area += 1
+    return area

CelerisLab/utils.py

@@ -0,0 +1,256 @@
+# CelerisLab/utils.py
+
+import pycuda.driver as cuda
+import subprocess
+import json
+
+from typing import NamedTuple, Optional, List, Tuple, Union
+
+
+class CudaDeviceInfo(NamedTuple):
+    name: str
+    compute_capability: str
+    multiprocessors: int
+    total_global_memory: int
+    max_shared_memory_per_block: int
+    max_threads_per_block: int
+    max_blocks_per_multiprocessor: int
+    device_interconnect: Optional[str] = None
+
+
+class FlowFieldConfig(NamedTuple):
+    data_type: str
+    dimensionality: int
+    lattice: int
+    field_dim_in_U: Tuple[int, int, int]
+    viscosity: float
+    velocity: float
+    boundary_conditions: Tuple[str, str, str, str, str, str]
+
+
+class CudaConfig(NamedTuple):
+    multi_gpu: bool
+    gpu_connection: str
+    required_cuda_capability: str
+    threads_per_block: int
+    unit_dimensions: Tuple[int, int, int]
+
+
+def check_cuda_device_availability(device_id=0):
+    if cuda.Device.count() == 0:
+        raise RuntimeError("No CUDA device is available.")
+
+    if device_id < 0 or device_id >= cuda.Device.count():
+        raise ValueError(
+            f"Invalid device_id {device_id}. Must be between 0 and {cuda.Device.count() - 1}."
+        )
+
+    try:
+        subprocess.check_output(["nvidia-smi", "--version"])
+    except subprocess.CalledProcessError:
+        raise RuntimeError("nvidia-smi is not available or not installed correctly.")
+
+
+def query_cuda_device_info(device_id=0) -> CudaDeviceInfo:
+    check_cuda_device_availability(device_id)
+
+    try:
+        output = subprocess.check_output(
+            ["nvidia-smi", "-q", "-d", "TOPOLOGY", "-i", str(device_id)], text=True
+        )
+        if "NVLink" in output:
+            device_interconnect = "NVLink"
+        elif "PCIe" in output:
+            device_interconnect = "PCIe"
+        else:
+            device_interconnect = "Unknown"
+    except Exception as e:
+        device_interconnect = None
+
+    device = cuda.Device(device_id)
+
+    return CudaDeviceInfo(
+        name=device.name(),
+        compute_capability=f"{device.compute_capability()[0]}.{device.compute_capability()[1]}",
+        multiprocessors=device.get_attribute(
+            cuda.device_attribute.MULTIPROCESSOR_COUNT
+        ),
+        total_global_memory=device.total_memory(),
+        max_shared_memory_per_block=device.get_attribute(
+            cuda.device_attribute.MAX_SHARED_MEMORY_PER_BLOCK
+        ),
+        max_threads_per_block=device.get_attribute(
+            cuda.device_attribute.MAX_THREADS_PER_BLOCK
+        ),
+        max_blocks_per_multiprocessor=device.get_attribute(
+            cuda.device_attribute.MAX_BLOCKS_PER_MULTIPROCESSOR
+        ),
+        device_interconnect=device_interconnect,
+    )
+
+
+def load_flow_field_config(config_path: str) -> FlowFieldConfig:
+    try:
+        with open(config_path, "r") as file:
+            config = json.load(file)
+
+        required_keys = [
+            "data_type",
+            "dimensionality",
+            "lattice",
+            "field_dim_in_U",
+            "viscosity",
+            "boundary_conditions",
+        ]
+        if not all(key in config for key in required_keys):
+            raise ValueError("Missing required configuration items.")
+
+        if config["data_type"] not in ["FP32", "FP64"]:
+            raise ValueError("Data type must be either FP32 or FP64.")
+
+        if config["dimensionality"] not in [2, 3]:
+            raise ValueError("Dimensionality must be either 2 or 3.")
+
+        if config["dimensionality"] == 2 and config["field_dim_in_U"][2] != 1:
+            raise ValueError(
+                "Field dimensions must be 1 in the third dimension for 2D simulations."
+            )
+
+        if config["lattice"] not in [9]:
+            raise ValueError("Lattice must be either 9 or 19.")
+
+        boundary_conditions = tuple(
+            condition
+            for key in ["x", "y", "z"]
+            for condition in config["boundary_conditions"].get(key, [])
+        )
+        if len(boundary_conditions) != 6:
+            raise ValueError("Boundary conditions must contain exactly six elements.")
+
+        return FlowFieldConfig(
+            data_type=config["data_type"],
+            dimensionality=config["dimensionality"],
+            lattice=config["lattice"],
+            field_dim_in_U=tuple(config["field_dim_in_U"]),
+            viscosity=config["viscosity"],
+            velocity=config["velocity"],
+            boundary_conditions=boundary_conditions,
+        )
+    except Exception as e:
+        raise RuntimeError(f"Failed to load or parse the flow field configuration: {e}")
+
+
+def load_cuda_config(config_path: str) -> CudaConfig:
+    try:
+        with open(config_path, "r") as file:
+            config = json.load(file)
+
+        required_keys = [
+            "multi_gpu",
+            "gpu_connection",
+            "required_cuda_capability",
+            "threads_per_block",
+            "X_1U",
+            "Y_1U",
+            "Z_1U",
+        ]
+
+        if not all(key in config for key in required_keys):
+            raise ValueError("Missing required configuration items.")
+
+        return CudaConfig(
+            multi_gpu=config["multi_gpu"],
+            gpu_connection=config["gpu_connection"],
+            required_cuda_capability=config["required_cuda_capability"],
+            threads_per_block=config["threads_per_block"],
+            unit_dimensions=(config["X_1U"], config["Y_1U"], config["Z_1U"]),
+        )
+    except Exception as e:
+        raise RuntimeError(f"Failed to load or parse the CUDA configuration: {e}")
+
+
+def check_cuda_capability(
+    field_config: FlowFieldConfig,
+    cuda_config: CudaConfig,
+    device_id: Union[int, List[int]] = None,
+):
+    SAFE_FACTOR = 0.8
+
+    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.")
+        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
+        device_info = query_cuda_device_info(device_id)
+
+        if device_info.compute_capability != cuda_config.required_cuda_capability:
+            raise ValueError(
+                f"Device {device_info.name} has compute capability {device_info.compute_capability}, but {cuda_config.required_cuda_capability} is required."
+            )
+
+        field_size = sum(
+            size * unit
+            for size, unit in zip(
+                field_config.field_dim_in_U, cuda_config.unit_dimensions
+            )
+        )
+        if (
+            device_info.total_global_memory * SAFE_FACTOR
+            < calc_field_memory_consumption(
+                field_size,
+                field_config.dimensionality,
+                field_config.lattice,
+                field_config.data_type,
+            )
+        ):
+            raise ValueError(
+                f"Device {device_info.name} does not have enough memory to store the flow field."
+            )
+
+        if (
+            device_info.max_threads_per_block * SAFE_FACTOR
+            < cuda_config.threads_per_block
+        ):
+            raise ValueError(
+                f"Device {device_info.name} does not have enough threads per block to run the simulation."
+            )
+
+        block_size = cuda_config.threads_per_block
+        if (
+            device_info.max_shared_memory_per_block * SAFE_FACTOR
+            < 2
+            * calc_field_memory_consumption(
+                block_size,
+                field_config.dimensionality,
+                field_config.lattice,
+                field_config.data_type,
+            )
+        ):
+            raise ValueError(
+                f"Device {device_info.name} does not have enough shared memory per block to run the simulation."
+            )
+
+
+def calc_field_memory_consumption(
+    field_size: int, dimensionality: int, directions: int, data_type: str
+) -> int:
+    if data_type == "FP32":
+        data_size = 4
+    elif data_type == "FP64":
+        data_size = 8
+    else:
+        raise ValueError(f"Unsupported data type {data_type}.")
+
+    return (
+        field_size * directions * data_size * 2
+        + field_size * dimensionality * data_size
+        + field_size
+    )

SB3/env_pinball.py

@@ -1,74 +0,0 @@
-import gymnasium as gym
-import numpy as np
-from gymnasium import spaces
-import ctypes
-from collections import deque
-
-lbm = ctypes.cdll.LoadLibrary('./lbm_sens.so')
-
-S_DIM, A_DIM = 6, 3
-action_amp = 5
-action_weight = 0.5
-sample_interval = 200
-max_steps = 320
-
-class CustomEnv(gym.Env):
-    """Custom Environment that follows gym interface."""
-
-    metadata = {"render_modes": ["human"], "render_fps": 1000/sample_interval}
-
-    def __init__(self, devicenum=0, Ccost=0.2):
-        super().__init__()
-        self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
-        self.observation_space = spaces.Box(low=-5, high=5, shape=(6,), dtype=np.float32)
-        self.fifo_rewards = deque(maxlen=50)
-        lbm.SetDevice(devicenum)
-        lbm.InitAll()
-        lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
-        lbm.CoreSolver.restype  = ctypes.POINTER(ctypes.c_float)
-        self.temps_init = lbm.CoreSolver(100*1000, 0.0, 0.0, 0.0, 0.0)
-        self.s = np.array([0.0] * S_DIM, dtype=np.float32)
-        for i in range(S_DIM):
-            self.s[i] = self.temps_init[i]
-        lbm.InitCPUMemory()
-        self.max_steps = max_steps
-        self.current_step = 0
-        self.Ccost = Ccost
-
-    def step(self, action):
-        assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
-        lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
-        lbm.CoreSolver.restype  = ctypes.POINTER(ctypes.c_float)
-        action = action_amp * action
-        temps = lbm.CoreSolver(sample_interval, action_weight, action[0], action[1], action[2])
-        for i in range(S_DIM):
-            self.s[i] = temps[i]
-        observation = np.hstack(self.s)
-        cd = self.s[0]+self.s[2]+self.s[4]
-        cl = self.s[1]+self.s[3]+self.s[5]
-        reward = float((1-self.Ccost)*np.exp(-np.abs(cd)/3)+self.Ccost*np.exp(-np.abs(cl)/3))
-        self.fifo_rewards.append(reward)
-        terminated = bool(np.mean(self.fifo_rewards) > 0.9)
-        truncated = bool(np.any(self.s > 3) or np.any(self.s < -3))
-        self.current_step += 1
-        if self.current_step >= self.max_steps:
-            terminated = True
-        info = {}
-        return observation, reward, terminated, truncated, info
-
-    def reset(self, seed=None, Ccost=0.2):
-        lbm.ResetAll()
-        for i in range(S_DIM):
-            self.s[i] = self.temps_init[i]
-        observation = np.hstack(self.s)
-        info = {}
-        self.current_step = 0
-        self.Ccost = Ccost
-        return observation, info
-
-    def render(self, episode=0, numstep=0):
-        lbm.OutputFlow.argtypes = (ctypes.c_int, ctypes.c_int, ctypes.c_int)
-        lbm.OutputFlow(episode, numstep, sample_interval)
-
-    def close(self):
-        lbm.Finalize()

SB3/jupyter.ipynb

@@ -1,148 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Using cuda device\n",
-      "Logging to ./tensorboard/PPO_1\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Process ForkServerProcess-2:\n",
-      "Process ForkServerProcess-1:\n",
-      "Traceback (most recent call last):\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
-      "    self.run()\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
-      "    self._target(*self._args, **self._kwargs)\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
-      "    observation, reward, terminated, truncated, info = env.step(data)\n",
-      "ValueError: not enough values to unpack (expected 5, got 4)\n",
-      "Traceback (most recent call last):\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
-      "    self.run()\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
-      "    self._target(*self._args, **self._kwargs)\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
-      "    observation, reward, terminated, truncated, info = env.step(data)\n",
-      "ValueError: not enough values to unpack (expected 5, got 4)\n",
-      "Process ForkServerProcess-4:\n",
-      "Traceback (most recent call last):\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
-      "    self.run()\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
-      "    self._target(*self._args, **self._kwargs)\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
-      "    observation, reward, terminated, truncated, info = env.step(data)\n",
-      "ValueError: not enough values to unpack (expected 5, got 4)\n",
-      "Process ForkServerProcess-3:\n",
-      "Traceback (most recent call last):\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 314, in _bootstrap\n",
-      "    self.run()\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/process.py\", line 108, in run\n",
-      "    self._target(*self._args, **self._kwargs)\n",
-      "  File \"/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py\", line 35, in _worker\n",
-      "    observation, reward, terminated, truncated, info = env.step(data)\n",
-      "ValueError: not enough values to unpack (expected 5, got 4)\n"
-     ]
-    },
-    {
-     "ename": "EOFError",
-     "evalue": "",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mEOFError\u001b[0m                                  Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[1], line 24\u001b[0m\n\u001b[1;32m     16\u001b[0m vec_env \u001b[38;5;241m=\u001b[39m SubprocVecEnv(env_fns)\n\u001b[1;32m     18\u001b[0m model \u001b[38;5;241m=\u001b[39m PPO(\n\u001b[1;32m     19\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMlpPolicy\u001b[39m\u001b[38;5;124m\"\u001b[39m, \n\u001b[1;32m     20\u001b[0m     env\u001b[38;5;241m=\u001b[39mvec_env, \n\u001b[1;32m     21\u001b[0m     n_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m64\u001b[39m,\n\u001b[1;32m     22\u001b[0m     tensorboard_log\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m./tensorboard/\u001b[39m\u001b[38;5;124m\"\u001b[39m, \n\u001b[1;32m     23\u001b[0m     verbose\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m---> 24\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m128\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m1000\u001b[39;49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/ppo/ppo.py:315\u001b[0m, in \u001b[0;36mPPO.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m    306\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mlearn\u001b[39m(\n\u001b[1;32m    307\u001b[0m     \u001b[38;5;28mself\u001b[39m: SelfPPO,\n\u001b[1;32m    308\u001b[0m     total_timesteps: \u001b[38;5;28mint\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    313\u001b[0m     progress_bar: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m    314\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m SelfPPO:\n\u001b[0;32m--> 315\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    316\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtotal_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    317\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcallback\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    318\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlog_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlog_interval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    319\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtb_log_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    320\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    321\u001b[0m \u001b[43m        \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    322\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/on_policy_algorithm.py:277\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m    274\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39menv \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    276\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m<\u001b[39m total_timesteps:\n\u001b[0;32m--> 277\u001b[0m     continue_training \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollect_rollouts\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrollout_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_rollout_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mn_steps\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    279\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m continue_training:\n\u001b[1;32m    280\u001b[0m         \u001b[38;5;28;01mbreak\u001b[39;00m\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/on_policy_algorithm.py:194\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.collect_rollouts\u001b[0;34m(self, env, callback, rollout_buffer, n_rollout_steps)\u001b[0m\n\u001b[1;32m    189\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    190\u001b[0m         \u001b[38;5;66;03m# Otherwise, clip the actions to avoid out of bound error\u001b[39;00m\n\u001b[1;32m    191\u001b[0m         \u001b[38;5;66;03m# as we are sampling from an unbounded Gaussian distribution\u001b[39;00m\n\u001b[1;32m    192\u001b[0m         clipped_actions \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mclip(actions, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mlow, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mhigh)\n\u001b[0;32m--> 194\u001b[0m new_obs, rewards, dones, infos \u001b[38;5;241m=\u001b[39m \u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43mclipped_actions\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    196\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39mnum_envs\n\u001b[1;32m    198\u001b[0m \u001b[38;5;66;03m# Give access to local variables\u001b[39;00m\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/base_vec_env.py:206\u001b[0m, in \u001b[0;36mVecEnv.step\u001b[0;34m(self, actions)\u001b[0m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \u001b[38;5;124;03mStep the environments with the given action\u001b[39;00m\n\u001b[1;32m    201\u001b[0m \n\u001b[1;32m    202\u001b[0m \u001b[38;5;124;03m:param actions: the action\u001b[39;00m\n\u001b[1;32m    203\u001b[0m \u001b[38;5;124;03m:return: observation, reward, done, information\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    205\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstep_async(actions)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstep_wait\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py:129\u001b[0m, in \u001b[0;36mSubprocVecEnv.step_wait\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep_wait\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m VecEnvStepReturn:\n\u001b[0;32m--> 129\u001b[0m     results \u001b[38;5;241m=\u001b[39m [remote\u001b[38;5;241m.\u001b[39mrecv() \u001b[38;5;28;01mfor\u001b[39;00m remote \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremotes]\n\u001b[1;32m    130\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwaiting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m    131\u001b[0m     obs, rews, dones, infos, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreset_infos \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39mresults)  \u001b[38;5;66;03m# type: ignore[assignment]\u001b[39;00m\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/subproc_vec_env.py:129\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstep_wait\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m VecEnvStepReturn:\n\u001b[0;32m--> 129\u001b[0m     results \u001b[38;5;241m=\u001b[39m [\u001b[43mremote\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrecv\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m remote \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremotes]\n\u001b[1;32m    130\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mwaiting \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m    131\u001b[0m     obs, rews, dones, infos, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreset_infos \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39mresults)  \u001b[38;5;66;03m# type: ignore[assignment]\u001b[39;00m\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:250\u001b[0m, in \u001b[0;36m_ConnectionBase.recv\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    248\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_closed()\n\u001b[1;32m    249\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_readable()\n\u001b[0;32m--> 250\u001b[0m buf \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_recv_bytes\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    251\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _ForkingPickler\u001b[38;5;241m.\u001b[39mloads(buf\u001b[38;5;241m.\u001b[39mgetbuffer())\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:414\u001b[0m, in \u001b[0;36mConnection._recv_bytes\u001b[0;34m(self, maxsize)\u001b[0m\n\u001b[1;32m    413\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_recv_bytes\u001b[39m(\u001b[38;5;28mself\u001b[39m, maxsize\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 414\u001b[0m     buf \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_recv\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    415\u001b[0m     size, \u001b[38;5;241m=\u001b[39m struct\u001b[38;5;241m.\u001b[39munpack(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m!i\u001b[39m\u001b[38;5;124m\"\u001b[39m, buf\u001b[38;5;241m.\u001b[39mgetvalue())\n\u001b[1;32m    416\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m size \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m:\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/multiprocessing/connection.py:383\u001b[0m, in \u001b[0;36mConnection._recv\u001b[0;34m(self, size, read)\u001b[0m\n\u001b[1;32m    381\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m n \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m    382\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m remaining \u001b[38;5;241m==\u001b[39m size:\n\u001b[0;32m--> 383\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mEOFError\u001b[39;00m\n\u001b[1;32m    384\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    385\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgot end of file during message\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
-      "\u001b[0;31mEOFError\u001b[0m: "
-     ]
-    }
-   ],
-   "source": [
-    "import os\n",
-    "os.environ['MKL_THREADING_LAYER'] = 'GNU'\n",
-    "import numpy as np\n",
-    "import gymnasium as gym\n",
-    "from env_pinball import CustomEnv\n",
-    "from stable_baselines3 import PPO\n",
-    "from stable_baselines3.common.vec_env import SubprocVecEnv\n",
-    "\n",
-    "def make_env(gpu_id):\n",
-    "    def _init():\n",
-    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)\n",
-    "        return CustomEnv(devicenum=gpu_id)\n",
-    "    return _init\n",
-    "\n",
-    "env_fns = [make_env(i) for i in range(4)]\n",
-    "vec_env = SubprocVecEnv(env_fns)\n",
-    "\n",
-    "model = PPO(\n",
-    "    \"MlpPolicy\", \n",
-    "    env=vec_env, \n",
-    "    n_steps=64,\n",
-    "    tensorboard_log=\"./tensorboard/\", \n",
-    "    verbose=1)\n",
-    "model.learn(total_timesteps=64*1000)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "vec_env = model.get_env()\n",
-    "obs = vec_env.reset()\n",
-    "\n",
-    "n_steps = 0\n",
-    "list_reward = {}\n",
-    "terminated = False\n",
-    "truncated = False\n",
-    "while n_steps < 500 and not terminated and not truncated:\n",
-    "    n_steps += 1\n",
-    "    action, _states = model.predict(observation=obs)\n",
-    "    obs, rewards, dones, info = vec_env.step(action)\n",
-    "    list_reward[n_steps] = rewards"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "pycuda_3_10",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.10.13"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

SB3/nohup.py

@@ -1,216 +0,0 @@
-# %%
-#!/usr/bin/env python3
-from env_pinball import CustomEnv
-import os
-import pickle
-import random
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.distributions.normal import Normal
-from torch.utils.tensorboard import SummaryWriter
-
-env = CustomEnv(devicenum=3)
-writer = SummaryWriter(log_dir='./tensorboard/DRL')
-
-# %%
-class Policy_Network(nn.Module):
-    """Parametrized Policy Network."""
-
-    def __init__(self, obs_space_dims: int, action_space_dims: int):
-        """Initializes a neural network that estimates the mean and standard deviation
-         of a normal distribution from which an action is sampled from.
-
-        Args:
-            obs_space_dims: Dimension of the observation space
-            action_space_dims: Dimension of the action space
-        """
-        super().__init__()
-
-        hidden_space1 = 256  # Nothing special with 16, feel free to change
-        hidden_space2 = 256  # Nothing special with 32, feel free to change
-
-        # Shared Network
-        self.shared_net = nn.Sequential(
-            nn.Linear(obs_space_dims, hidden_space1),
-            nn.Tanh(),
-            nn.Linear(hidden_space1, hidden_space2),
-            nn.Tanh(),
-        )
-
-        # Policy Mean specific Linear Layer
-        self.policy_mean_net = nn.Sequential(
-            nn.Linear(hidden_space2, action_space_dims)
-        )
-
-        # Policy Std Dev specific Linear Layer
-        self.policy_stddev_net = nn.Sequential(
-            nn.Linear(hidden_space2, action_space_dims)
-        )
-
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        """Conditioned on the observation, returns the mean and standard deviation
-         of a normal distribution from which an action is sampled from.
-
-        Args:
-            x: Observation from the environment
-
-        Returns:
-            action_means: predicted mean of the normal distribution
-            action_stddevs: predicted standard deviation of the normal distribution
-        """
-        shared_features = self.shared_net(x.float())
-
-        action_means = self.policy_mean_net(shared_features)
-        action_means = torch.tanh(action_means)
-
-        action_stddevs = torch.log(
-            1 + torch.exp(self.policy_stddev_net(shared_features))
-        )
-
-        return action_means, action_stddevs
-
-# %%
-class REINFORCE:
-    """REINFORCE algorithm."""
-
-    def __init__(self, obs_space_dims: int, action_space_dims: int):
-        """Initializes an agent that learns a policy via REINFORCE algorithm [1]
-        to solve the task at hand (Inverted Pendulum v4).
-
-        Args:
-            obs_space_dims: Dimension of the observation space
-            action_space_dims: Dimension of the action space
-        """
-
-        # Hyperparameters
-        self.learning_rate = 1e-4  # Learning rate for policy optimization
-        self.gamma = 0.99  # Discount factor
-        self.eps = 1e-6  # small number for mathematical stability
-
-        self.probs = []  # Stores probability values of the sampled action
-        self.rewards = []  # Stores the corresponding rewards
-
-        self.net = Policy_Network(obs_space_dims, action_space_dims)
-        self.optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.learning_rate)
-
-    def sample_action(self, state: np.ndarray) -> float:
-        """Returns an action, conditioned on the policy and observation.
-
-        Args:
-            state: Observation from the environment
-
-        Returns:
-            action: Action to be performed
-        """
-        state = torch.tensor(np.array([state]))
-        action_means, action_stddevs = self.net(state)
-
-        # create a normal distribution from the predicted
-        #   mean and standard deviation and sample an action
-        distrib = Normal(action_means[0] + self.eps, action_stddevs[0] + self.eps)
-        action = distrib.sample()
-        prob = distrib.log_prob(action)
-
-        action = torch.tanh(action)
-        action = action.numpy()
-
-        self.probs.append(prob)
-
-        return action
-
-    def update(self):
-        """Updates the policy network's weights."""
-        running_g = 0
-        gs = []
-
-        # Discounted return (backwards) - [::-1] will return an array in reverse
-        for R in self.rewards[::-1]:
-            running_g = R + self.gamma * running_g
-            gs.insert(0, running_g)
-
-        deltas = torch.tensor(gs)
-
-        loss = 0
-        # minimize -1 * prob * reward obtained
-        for log_prob, delta in zip(self.probs, deltas):
-            loss += log_prob.mean() * delta * (-1)
-
-        # Update the policy network
-        self.optimizer.zero_grad()
-        loss.backward()
-        self.optimizer.step()
-
-        # Empty / zero out all episode-centric/related variables
-        self.probs = []
-        self.rewards = []
-
-# %%
-total_num_episodes = int(5e3)  # Total number of episodes
-obs_space_dims = 6
-action_space_dims = 3
-rewards_over_seeds = []
-MAX_REWARD = 0
-
-# Check if there is a saved state
-if os.path.exists('saved_state.pkl'):
-    with open('saved_state.pkl', 'rb') as f:
-        i_seed, episode, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD = pickle.load(f)
-    os.remove('saved_state.pkl')  # Remove the saved state
-else:
-    i_seed = 0
-    episode = 0
-    agent = None
-    reward_over_episodes = None
-
-for seed in [1][i_seed:]:  # Fibonacci seeds
-    # set seed
-    torch.manual_seed(seed)
-    random.seed(seed)
-    np.random.seed(seed)
-
-    # Reinitialize agent every seed
-    if agent is None or reward_over_episodes is None:
-        agent = REINFORCE(obs_space_dims, action_space_dims)
-        reward_over_episodes = []
-
-    while episode < total_num_episodes+1:
-        obs, info = env.reset(Ccost=0.2+episode/total_num_episodes*0.6)
-        steps = 0
-        done = False
-        terminated = False
-        truncated = False
-        reward_over_steps = []
-        while not done:
-            action = agent.sample_action(obs)
-            obs, reward, terminated, truncated, info = env.step(action)
-            agent.rewards.append(reward)
-            reward_over_steps.append(reward)
-            steps += 1
-            done = terminated or truncated
-
-        avg_reward = np.mean(reward_over_steps[-64:])
-        reward_over_episodes.append(np.array([avg_reward], dtype=np.float32))
-        agent.update()
-
-        if episode % 10 == 0:
-            # print("Episode:", episode, "Average Reward:", int(avg_reward))
-            writer.add_scalar('Average Reward', int(avg_reward), episode)
-
-        if avg_reward > MAX_REWARD:
-            MAX_REWARD = avg_reward
-            with open('saved_model_'+str(seed)+'.pkl', 'wb') as f:
-                pickle.dump((episode + 1, agent, reward_over_episodes, MAX_REWARD), f)
-
-        # Save the current state at the end of each episode
-        with open('saved_state.pkl', 'wb') as f:
-            pickle.dump((i_seed, episode + 1, agent, reward_over_episodes, rewards_over_seeds, MAX_REWARD), f)
-
-        episode += 1
-    episode = 0
-    MAX_REWARD = 0
-    i_seed += 1
-    rewards_over_seeds.append(reward_over_episodes)
-    agent = None  # Reset the agent
-    reward_over_episodes = None  # Reset the reward_over_episodes
-# %%

configs/config_cuda.json

@@ -0,0 +1,9 @@
+{
+    "multi_gpu": false,
+    "gpu_connection": "NVLink",
+    "required_cuda_capability": "7.0",
+    "threads_per_block": 128,
+    "X_1U": 128,
+    "Y_1U": 32,
+    "Z_1U": 1
+}

configs/config_flowfield.json

@@ -0,0 +1,13 @@
+{
+    "data_type": "FP32",
+    "dimensionality": 2,
+    "lattice": 9,
+    "field_dim_in_U": [10, 16, 1],
+    "viscosity": 0.004,
+    "velocity": 0.01,
+    "boundary_conditions": {
+        "x": ["parabolic", "outflow"],
+        "y": ["noslip", "noslip"],
+        "z": ["none", "none"]
+    }
+}

configs/config_gym.json

@@ -0,0 +1,3 @@
+{
+    
+}

lbm_kernel.py

@@ -1,9 +0,0 @@
-import pycuda.driver as cuda
-import pycuda.autoinit
-from pycuda.compiler import SourceModule
-import numpy as np
-
-with open('./cuda/kernel.cu', 'r') as file_k:
-    code = file_k.read()
-
-kernel = SourceModule(code)

scripts/d1a3o12.py

@@ -0,0 +1,57 @@
+import os
+os.environ['MKL_THREADING_LAYER'] = 'GNU'
+os.environ["OMP_NUM_THREADS"] = "8"
+os.environ["MKL_NUM_THREADS"] = "8"
+import torch
+import numpy as np
+from torch.nn import Module
+import gymnasium as gym
+from gym_env import CustomEnv
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import SubprocVecEnv
+from stable_baselines3.common.vec_env import DummyVecEnv
+from sb3_contrib import RecurrentPPO
+from torch.utils.tensorboard import SummaryWriter
+
+current_dir = os.path.dirname(os.path.abspath("__file__"))
+parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
+
+class Sin(Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.sin(x)
+
+if __name__ == '__main__':
+    
+    vec_env = CustomEnv(device_id=1)
+    name = "d1a3o12_c1"
+
+    model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_a0"), env=vec_env, device=torch.device("cuda:1"))
+
+    # model = PPO(
+    #     "MlpPolicy", 
+    #     policy_kwargs=dict(activation_fn=Sin),
+    #     env=vec_env, 
+    #     device=torch.device("cuda:1"),
+    #     verbose=0)
+
+    writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
+    max_reward = 0
+
+    for i in range(100):
+        model.learn(total_timesteps=480)
+        test_env = model.get_env()
+        test_obs = test_env.reset()
+        list_reward = []
+        for step in range(480):
+            test_action, _states = model.predict(observation=test_obs)
+            test_obs, test_rewards, test_dones, info = test_env.step(test_action)
+            list_reward.append(test_rewards)
+
+        avg_reward = np.mean(list_reward[-240:])
+        writer.add_scalar('Reward', np.mean(avg_reward), i)
+        if avg_reward > max_reward:
+            max_reward = avg_reward
+            model.save(os.path.join(parent_dir, "models", name + ".zip"))

scripts/gym_env.py

@@ -0,0 +1,198 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium import spaces
+import ctypes
+from collections import deque
+from typing import Tuple
+import sys
+import os
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ProcessPoolExecutor
+import queue
+
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+
+current_dir = os.path.dirname(os.path.abspath("__file__"))
+parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
+sys.path.append(parent_dir)
+from CelerisLab import FlowField
+from CelerisLab import utils
+
+config_cuda = utils.load_cuda_config(
+    os.path.join(parent_dir, "configs", "config_cuda.json")
+)
+config_field = utils.load_flow_field_config(
+    os.path.join(parent_dir, "configs", "config_flowfield.json")
+)
+
+S_DIM, A_DIM = 12, 3
+U0 = config_field.velocity
+T0 = 1000
+SAMPLE_INTERVAL = 800
+FIFO_LEN = 120
+CONV_LEN = 60
+MAX_STEPS = 640
+if config_field.data_type == "FP32":
+    DATA_TYPE = np.float32
+else:
+    raise ValueError(f"Unsupported data type {config_field.data_type}.")
+
+
+class CustomEnv(gym.Env):
+    """Custom Environment that follows gym interface."""
+
+    metadata = {"render_modes": ["human"], "render_fps": T0 / SAMPLE_INTERVAL}
+
+    def __init__(self, device_id=0):
+        super().__init__()
+        self.action_space = spaces.Box(low=-1, high=1, shape=(A_DIM,), dtype=DATA_TYPE)
+        self.observation_space = spaces.Box(
+            low=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
+        )
+        self.fifo_states = deque(maxlen=FIFO_LEN)
+        self.target_states = np.empty((0, 6), dtype=DATA_TYPE)
+        self.force_norm_fact = 1.0
+        self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
+        self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)
+
+        self.flow_field = FlowField(config_field, config_cuda, device_id)
+        L0 = 20
+        U0 = config_field.velocity
+        NX = self.flow_field.FIELD_SHAPE[0]
+        NY = self.flow_field.FIELD_SHAPE[1]
+        center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)
+        self.flow_field.add_cylinder(center, L0)
+        center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0)
+        self.flow_field.add_sensor(center, L0 / 4)
+        center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0)
+        self.flow_field.add_sensor(center, L0 / 4)
+        center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 - 2 * L0, 0)
+        self.flow_field.add_sensor(center, L0 / 4)
+        self.flow_field.run(int(4*NX/U0), np.zeros(4, dtype=DATA_TYPE))
+
+        for i in range(FIFO_LEN):
+            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(4, dtype=DATA_TYPE))
+            new_state = self.flow_field.obs.copy()[2:8]
+            self.target_states = np.vstack((self.target_states, new_state))
+
+        self.flow_field.apply_ddf()
+        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        self.flow_field.run(int(4*NX/U0), np.zeros(7, dtype=DATA_TYPE))
+        self.flow_field.get_ddf()
+
+        for i in range(FIFO_LEN):
+            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(7, dtype=DATA_TYPE))
+            self.fifo_states.append(self.flow_field.obs.copy()[2:14])
+
+        temp_states = np.array(self.fifo_states)
+        self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 6:12]))
+        for i in range(6):
+            self.sens_deviation[i] = np.mean(temp_states[:, i])
+            self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i] - self.sens_deviation[i]))
+            self.target_states[:, i] = (self.target_states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i]
+
+
+    def step(self, action):
+        assert self.action_space.contains(action), "%r (%s) invalid" % (
+            action,
+            type(action),
+        )
+
+        # barrier = threading.Barrier(2)
+        result_queue = queue.Queue()
+
+        def run_flow_field(action):
+            self.flow_field.context.push()
+            U0 = config_field.velocity
+            try:
+                temp = np.zeros(7, dtype=DATA_TYPE)
+                temp[4:7] = np.array((action*8+[0,-4,4])*U0, dtype=DATA_TYPE)
+                self.flow_field.run(SAMPLE_INTERVAL, temp)
+            finally:
+                self.flow_field.context.pop()
+            # barrier.wait()
+            self.fifo_states.append(self.flow_field.obs.copy()[2:14])
+
+        def calc_reward():
+            states = np.array(self.fifo_states)
+            forces = states[-1, 6:12] / self.force_norm_fact
+            cd = (forces[0] + forces[2] + forces[4]) / 3
+            cl = (forces[1] + forces[3] + forces[5]) / 3
+            sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact
+
+            similarities = 0.0
+
+            def calc_lag(target, state):
+                target_mean = np.mean(target)
+                state_mean = np.mean(state)
+
+                correlation = np.correlate(target - target_mean, state - state_mean, "full")
+                lags = np.arange(-len(target) + 1, len(target))
+                max_lag = lags[np.argmax(correlation)]
+                return max_lag
+                
+            def calc_sim(target, state, lag):
+                target_mean = np.mean(target)
+                state_mean = np.mean(state)
+                target_std = np.std(target)
+
+                aligned_state = np.roll(state, lag)
+
+                if lag >= 0:
+                    seq_target = target[-CONV_LEN:]-target_mean
+                    seq_state = aligned_state[-CONV_LEN:]-state_mean
+                else:
+                    seq_target = target[:CONV_LEN]-target_mean
+                    seq_state = aligned_state[:CONV_LEN]-state_mean
+
+                seq_diff = seq_target - seq_state
+                sim_cor = 10*(np.corrcoef(seq_target, seq_state)[0, 1] - 1)
+                sim_div = -np.abs((target_mean - state_mean) / target_std * 0.75)
+                sim_amp = -np.abs(np.std(seq_diff) / target_std * 2)
+
+                return np.exp((sim_cor + sim_div + sim_amp) / 3)
+
+            id_sens = 0
+            target_seq = self.target_states[:, id_sens]
+            state_seq = (states[:, id_sens] - self.sens_deviation[id_sens]) / self.sens_norm_fact[id_sens]
+            lag = calc_lag(target_seq, state_seq)
+            similarities += calc_sim(target_seq, state_seq, lag) / 6
+
+            for i in range(1, 6):
+                target_seq = self.target_states[:, i]
+                state_seq = (states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i]
+                similarities += calc_sim(target_seq, state_seq, lag) / 6
+
+            reward_cd = np.exp(-np.abs(cd * 80))
+            reward_cl = np.exp(-np.abs(cl * 20))
+            # reward_sim = np.exp(2 * (similarities - 1))
+            reward_sim = similarities
+            reward = np.minimum(0.3 * reward_cd + 0.3 * reward_cl + 0.4 * reward_sim, 1.0)
+            # barrier.wait()
+            result_queue.put((np.hstack([forces, sens]), reward))
+
+        run_flow_field(action)
+        calc_reward()
+        observation, reward = result_queue.get()
+
+        truncated = bool(np.any(observation > 1) or np.any(observation < -1))
+        observation = np.clip(observation, -1, 1)
+        # truncated = False
+        return observation, float(reward), False, truncated, {}
+    
+    def reset(self, seed=None):
+        self.flow_field.apply_ddf()
+        return np.zeros(S_DIM, dtype=np.float32), {}
+    
+    def render(self, mode="human"):
+        pass
+
+    def close(self):
+        self.flow_field.__del__()

setup.py

@@ -0,0 +1,17 @@
+from setuptools import setup, find_packages
+
+setup(
+	name='CelerisLab',
+	version='0.1',
+	packages=find_packages(),
+	install_requires=[
+		'pycuda',
+		'numpy',
+        'json'
+	],
+	entry_points={
+		'console_scripts': [
+			'CelerisLab=CelerisLab.driver:main',
+		],
+	},
+)

test.cu

@@ -1,13 +0,0 @@
-#include <stdio.h>
-#include <cuda.h>
-
-extern "C" {
-__global__ void add(float *x, float *y)
-{
-  int index = threadIdx.x;
-  int stride = blockDim.x;
-
-  for (int i = index; i < 100; i += stride)
-    y[i] = x[i] + y[i];
-}
-}

test.ipynb

@@ -1,84 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "TypeError",
-     "evalue": "No registered converter was able to produce a C++ rvalue of type unsigned int from this Python object of type numpy.int64",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[1], line 25\u001b[0m\n\u001b[1;32m     22\u001b[0m drv\u001b[38;5;241m.\u001b[39mmemcpy_htod(y_gpu, y)\n\u001b[1;32m     24\u001b[0m \u001b[38;5;66;03m# 调用函数\u001b[39;00m\n\u001b[0;32m---> 25\u001b[0m \u001b[43madd_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_gpu\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my_gpu\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mblock\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrid\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mn\u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m256\u001b[39;49m\u001b[43m,\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     27\u001b[0m \u001b[38;5;66;03m# 将结果复制回CPU\u001b[39;00m\n\u001b[1;32m     28\u001b[0m drv\u001b[38;5;241m.\u001b[39mmemcpy_dtoh(y, y_gpu)\n",
-      "File \u001b[0;32m~/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/pycuda/driver.py:502\u001b[0m, in \u001b[0;36m_add_functionality.<locals>.function_call\u001b[0;34m(func, *args, **kwargs)\u001b[0m\n\u001b[1;32m    498\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtime\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m time\n\u001b[1;32m    500\u001b[0m     start_time \u001b[38;5;241m=\u001b[39m time()\n\u001b[0;32m--> 502\u001b[0m \u001b[43mfunc\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_launch_kernel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgrid\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mblock\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43marg_buf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshared\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    504\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m post_handlers \u001b[38;5;129;01mor\u001b[39;00m time_kernel:\n\u001b[1;32m    505\u001b[0m     Context\u001b[38;5;241m.\u001b[39msynchronize()\n",
-      "\u001b[0;31mTypeError\u001b[0m: No registered converter was able to produce a C++ rvalue of type unsigned int from this Python object of type numpy.int64"
-     ]
-    }
-   ],
-   "source": [
-    "import pycuda.autoinit\n",
-    "import pycuda.driver as drv\n",
-    "import numpy as np\n",
-    "\n",
-    "# 加载PTX文件\n",
-    "mod = drv.module_from_file(\"test.ptx\")\n",
-    "\n",
-    "# 获取函数\n",
-    "add_func = mod.get_function(\"add\")\n",
-    "\n",
-    "# 创建数据\n",
-    "n = np.uint32(100)  # Convert to unsigned int\n",
-    "x = np.random.rand(n).astype(np.float32)\n",
-    "y = np.random.rand(n).astype(np.float32)\n",
-    "\n",
-    "# 分配内存\n",
-    "x_gpu = drv.mem_alloc(x.nbytes)\n",
-    "y_gpu = drv.mem_alloc(y.nbytes)\n",
-    "\n",
-    "# 将数据复制到GPU\n",
-    "drv.memcpy_htod(x_gpu, x)\n",
-    "drv.memcpy_htod(y_gpu, y)\n",
-    "\n",
-    "# 调用函数\n",
-    "add_func(x_gpu, y_gpu, block=(256,1,1), grid=(n//256,1))\n",
-    "\n",
-    "# 将结果复制回CPU\n",
-    "drv.memcpy_dtoh(y, y_gpu)\n",
-    "\n",
-    "# 检查结果\n",
-    "print(y)"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "pycuda_3_10",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.10.13"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

test.ptx

@@ -1,104 +0,0 @@
-//
-// Generated by NVIDIA NVVM Compiler
-//
-// Compiler Build ID: CL-32688072
-// Cuda compilation tools, release 12.1, V12.1.105
-// Based on NVVM 7.0.1
-//
-
-.version 8.1
-.target sm_52
-.address_size 64
-
-	// .globl	add
-
-.visible .entry add(
-	.param .u64 add_param_0,
-	.param .u64 add_param_1
-)
-{
-	.reg .pred 	%p<6>;
-	.reg .f32 	%f<16>;
-	.reg .b32 	%r<22>;
-	.reg .b64 	%rd<25>;
-
-
-	ld.param.u64 	%rd11, [add_param_0];
-	ld.param.u64 	%rd12, [add_param_1];
-	cvta.to.global.u64 	%rd1, %rd12;
-	cvta.to.global.u64 	%rd2, %rd11;
-	mov.u32 	%r1, %ntid.x;
-	mov.u32 	%r20, %tid.x;
-	setp.gt.s32 	%p1, %r20, 99;
-	@%p1 bra 	$L__BB0_7;
-
-	mov.u32 	%r12, 99;
-	sub.s32 	%r13, %r12, %r20;
-	div.u32 	%r3, %r13, %r1;
-	add.s32 	%r14, %r3, 1;
-	and.b32  	%r19, %r14, 3;
-	setp.eq.s32 	%p2, %r19, 0;
-	@%p2 bra 	$L__BB0_4;
-
-	mul.wide.s32 	%rd13, %r20, 4;
-	add.s64 	%rd24, %rd1, %rd13;
-	mul.wide.s32 	%rd4, %r1, 4;
-	add.s64 	%rd23, %rd2, %rd13;
-
-$L__BB0_3:
-	.pragma "nounroll";
-	ld.global.f32 	%f1, [%rd24];
-	ld.global.f32 	%f2, [%rd23];
-	add.f32 	%f3, %f2, %f1;
-	st.global.f32 	[%rd24], %f3;
-	add.s32 	%r20, %r20, %r1;
-	add.s64 	%rd24, %rd24, %rd4;
-	add.s64 	%rd23, %rd23, %rd4;
-	add.s32 	%r19, %r19, -1;
-	setp.ne.s32 	%p3, %r19, 0;
-	@%p3 bra 	$L__BB0_3;
-
-$L__BB0_4:
-	setp.lt.u32 	%p4, %r3, 3;
-	@%p4 bra 	$L__BB0_7;
-
-	mul.wide.s32 	%rd10, %r1, 4;
-
-$L__BB0_6:
-	mul.wide.s32 	%rd14, %r20, 4;
-	add.s64 	%rd15, %rd2, %rd14;
-	add.s64 	%rd16, %rd1, %rd14;
-	ld.global.f32 	%f4, [%rd16];
-	ld.global.f32 	%f5, [%rd15];
-	add.f32 	%f6, %f5, %f4;
-	st.global.f32 	[%rd16], %f6;
-	add.s64 	%rd17, %rd15, %rd10;
-	add.s64 	%rd18, %rd16, %rd10;
-	ld.global.f32 	%f7, [%rd18];
-	ld.global.f32 	%f8, [%rd17];
-	add.f32 	%f9, %f8, %f7;
-	st.global.f32 	[%rd18], %f9;
-	add.s32 	%r15, %r20, %r1;
-	add.s32 	%r16, %r15, %r1;
-	add.s64 	%rd19, %rd17, %rd10;
-	add.s64 	%rd20, %rd18, %rd10;
-	ld.global.f32 	%f10, [%rd20];
-	ld.global.f32 	%f11, [%rd19];
-	add.f32 	%f12, %f11, %f10;
-	st.global.f32 	[%rd20], %f12;
-	add.s32 	%r17, %r16, %r1;
-	add.s64 	%rd21, %rd19, %rd10;
-	add.s64 	%rd22, %rd20, %rd10;
-	ld.global.f32 	%f13, [%rd22];
-	ld.global.f32 	%f14, [%rd21];
-	add.f32 	%f15, %f14, %f13;
-	st.global.f32 	[%rd22], %f15;
-	add.s32 	%r20, %r17, %r1;
-	setp.lt.s32 	%p5, %r20, 100;
-	@%p5 bra 	$L__BB0_6;
-
-$L__BB0_7:
-	ret;
-
-}
-