new branch test

CelerisLab/driver.py

@@ -2,7 +2,7 @@
 
 import pycuda.driver as cuda
 import numpy as np
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Optional
 
 from . import utils
 from . import preprocess as preproc
@@ -93,13 +93,13 @@ class FlowField:
         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.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.objects = {}
         self.action = np.zeros(0, dtype=self.DATA_TYPE)
@@ -118,7 +118,7 @@ class FlowField:
         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):
+    def add_cylinder(self, center: Tuple[float, float, float], radius: float, id_obj: Optional[int] = None):
         x_c, y_c, z_c = center
 
         if (
@@ -130,10 +130,13 @@ class FlowField:
             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: