Implemented 2d array for rapid block buffer (#71)

- added set_data_buffer_rapid_capture()
- added set_data_buffers_rapid_capture()

Author: JamesPicoTech

pypicosdk/base.py

@@ -1233,17 +1233,15 @@ def set_data_buffer_for_enabled_channels(
         # Rapid
         if captures > 0:
             for channel in self.range:
-                buffer = []
-                for capture_segment in range(captures):
-                    buffer.append(self.set_data_buffer(channel, samples, segment + capture_segment, datatype, ratio_mode, action=ACTION.ADD))
-                channels_buffer[channel] = buffer
+                np_buffer = self.set_data_buffer_rapid_capture(channel, samples, captures, segment, datatype, ratio_mode, action=ACTION.ADD)
+                channels_buffer[channel] = np_buffer
         # Single
         else:
             for channel in self.range:
                 channels_buffer[channel] = self.set_data_buffer(channel, samples, segment, datatype, ratio_mode, action=ACTION.ADD)
 
         return channels_buffer
-    
+        
     def set_data_buffer(
         self,
         channel,
@@ -1303,6 +1301,69 @@ def set_data_buffer(
         )
         return buffer
 
+        
+    def set_data_buffer_rapid_capture(
+            self,
+            channel,
+            samples,
+            captures,
+            segment=0,
+            datatype=DATA_TYPE.INT16_T,
+            ratio_mode=RATIO_MODE.RAW,
+            action=ACTION.CLEAR_ALL | ACTION.ADD,
+        ) -> np.ndarray | None:
+        """
+        Allocates and assigns multiple data buffers for rapid block capture on a specified channel.
+
+        Args:
+            channel (int): The channel to associate the buffer with (e.g., CHANNEL.A).
+            samples (int): Number of samples to allocate in the buffer.
+            captures (int): Number of rapid block captures
+            segment (int, optional): Memory segment to start at. 
+            datatype (DATA_TYPE, optional): C data type for the buffer (e.g., INT16_T). 
+            ratio_mode (RATIO_MODE, optional): Downsampling mode. 
+            action (ACTION, optional): Action to apply to the data buffer (e.g., CLEAR_ALL | ADD).
+
+        Returns:
+            np.array | None: The allocated buffer or ``None`` when clearing existing buffers.
+
+        Raises:
+            PicoSDKException: If an unsupported data type is provided.
+        """
+        if samples == 0:
+            buffer = None
+            buf_ptr = None
+        else:
+            # Map to NumPy dtype
+            dtype_map = {
+                DATA_TYPE.INT8_T: np.int8,
+                DATA_TYPE.INT16_T: np.int16,
+                DATA_TYPE.INT32_T: np.int32,
+                DATA_TYPE.INT64_T: np.int64,
+                DATA_TYPE.UINT32_T: np.uint32,
+            }
+
+            np_dtype = dtype_map.get(datatype)
+            if np_dtype is None:
+                raise PicoSDKException("Invalid datatype selected for buffer")
+
+            buffer = np.zeros((captures, samples), dtype=np_dtype)
+        
+        for i in range(captures):
+            self._call_attr_function(
+                "SetDataBuffer",
+                self.handle,
+                channel,
+                npc.as_ctypes(buffer[i]),
+                samples,
+                datatype,
+                segment + i,
+                ratio_mode,
+                action,
+            )
+
+        return buffer
+
     def set_data_buffers(
         self,
         channel,
@@ -1313,7 +1374,7 @@ def set_data_buffers(
         action=ACTION.CLEAR_ALL | ACTION.ADD,
     ) -> tuple[np.ndarray, np.ndarray]:
         """
-        Allocate and assign max and min NumPy-backed data buffers for 6000A series.
+        Allocate and assign max and min NumPy-backed data buffers.
 
         Args:
             channel (int): The channel to associate the buffers with.
@@ -1362,6 +1423,70 @@ def set_data_buffers(
         )
 
         return buffer_min, buffer_max
+    
+    def set_data_buffers_rapid_capture(
+            self,
+            channel,
+            samples,
+            captures,
+            segment=0,
+            datatype=DATA_TYPE.INT16_T,
+            ratio_mode=RATIO_MODE.RAW,
+            action=ACTION.CLEAR_ALL | ACTION.ADD,
+        ) -> np.ndarray | None:
+        """
+        Allocate and assign max and min NumPy-backed data buffers for rapid block 
+        capture on a specified channel.
+
+        Args:
+            channel (int): The channel to associate the buffer with (e.g., CHANNEL.A).
+            samples (int): Number of samples to allocate in the buffer.
+            captures (int): Number of rapid block captures
+            segment (int, optional): Memory segment to start at. 
+            datatype (DATA_TYPE, optional): C data type for the buffer (e.g., INT16_T). 
+            ratio_mode (RATIO_MODE, optional): Downsampling mode. 
+            action (ACTION, optional): Action to apply to the data buffer (e.g., CLEAR_ALL | ADD).
+
+        Returns:
+            np.array | None: The allocated buffer or ``None`` when clearing existing buffers.
+
+        Raises:
+            PicoSDKException: If an unsupported data type is provided.
+        """
+        if samples == 0:
+            buffer = None
+            buf_ptr = None
+        else:
+            # Map to NumPy dtype
+            dtype_map = {
+                DATA_TYPE.INT8_T: np.int8,
+                DATA_TYPE.INT16_T: np.int16,
+                DATA_TYPE.INT32_T: np.int32,
+                DATA_TYPE.INT64_T: np.int64,
+                DATA_TYPE.UINT32_T: np.uint32,
+            }
+
+            np_dtype = dtype_map.get(datatype)
+            if np_dtype is None:
+                raise PicoSDKException("Invalid datatype selected for buffer")
+
+            buffer = np.zeros((captures, samples, 2), dtype=np_dtype)
+        
+        for i in range(captures):
+            self._call_attr_function(
+                "SetDataBuffers",
+                self.handle,
+                channel,
+                npc.as_ctypes(buffer[i][0]),
+                npc.as_ctypes(buffer[i][1]),
+                samples,
+                datatype,
+                segment + i,
+                ratio_mode,
+                action,
+            )
+
+        return buffer
 
 
     # Run functions