modified read_n also for File and RawFile

python/CMakeLists.txt

@@ -28,8 +28,10 @@ target_link_libraries(_aare PRIVATE aare_core aare_compiler_flags)
 set( PYTHON_FILES
     aare/__init__.py
     aare/CtbRawFile.py
+    aare/RawFile.py
     aare/transform.py
     aare/ScanParameters.py
+    aare/utils.py
 )
 
 # Copy the python files to the build directory

python/aare/CtbRawFile.py

@@ -8,14 +8,15 @@ class CtbRawFile(_aare.CtbRawFile):
     
     Args:  
             fname (pathlib.Path | str): Path to the file to be read.
+            chunk_size (int): Number of frames to read at a time. Default is 1.
             transform (function): Function to apply to the data after reading it. 
                 The function should take a numpy array of type uint8 and return one
                 or several numpy arrays.
     """
-    def __init__(self, fname, transform = None, chunk_size = 1):
+    def __init__(self, fname, chunk_size = 1, transform = None):
         super().__init__(fname)
-        self.transform = transform
         self._chunk_size = chunk_size
+        self._transform = transform
 
 
     def read_frame(self, frame_index: int | None = None ) -> tuple:
@@ -44,8 +45,8 @@ def read_frame(self, frame_index: int | None = None ) -> tuple:
             header = header[0]
 
 
-        if self.transform:
-            res = self.transform(data)
+        if self._transform:
+            res = self._transform(data)
             if isinstance(res, tuple):
                 return header, *res
             else:

python/aare/RawFile.py

@@ -0,0 +1,66 @@
+from . import _aare
+import numpy as np
+from .ScanParameters import ScanParameters
+
+class RawFile(_aare.RawFile):
+    def __init__(self, fname, chunk_size = 1):
+        super().__init__(fname)
+        self._chunk_size = chunk_size
+
+
+    def read(self) -> tuple:
+        """Read the entire file.
+        Seeks to the beginning of the file before reading.
+
+        Returns:
+            tuple: header, data
+        """
+        self.seek(0)
+        return self.read_n(self.total_frames)
+
+    @property
+    def scan_parameters(self):
+        """Return the scan parameters.
+
+        Returns:
+            ScanParameters: Scan parameters.
+        """
+        return ScanParameters(self.master.scan_parameters)
+
+    @property
+    def master(self):
+        """Return the master file.
+
+        Returns:
+            RawMasterFile: Master file.
+        """
+        return super().master()
+
+    def __len__(self) -> int:
+        """Return the number of frames in the file.
+
+        Returns:
+            int: Number of frames in file.
+        """
+        return super().frames_in_file
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            if self._chunk_size == 1:
+                return self.read_frame()
+            else:
+                return self.read_n(self._chunk_size)
+
+
+        except RuntimeError:
+            # TODO! find a good way to check that we actually have the right exception
+            raise StopIteration

python/aare/__init__.py

@@ -2,10 +2,13 @@
 from . import _aare
 
 
-from ._aare import File, RawFile, RawMasterFile, RawSubFile
+from ._aare import File, RawMasterFile, RawSubFile
 from ._aare import Pedestal, ClusterFinder, VarClusterFinder
 from ._aare import DetectorType
 from ._aare import ClusterFile
 
 from .CtbRawFile import CtbRawFile
-from .ScanParameters import ScanParameters
+from .RawFile import RawFile
+from .ScanParameters import ScanParameters
+
+from .utils import random_pixels, random_pixel

python/aare/utils.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+def random_pixels(n_pixels, xmin=0, xmax=512, ymin=0, ymax=1024):
+    """Return a list of random pixels.
+    
+    Args:
+        n_pixels (int): Number of pixels to return.
+        rows (int): Number of rows in the image.
+        cols (int): Number of columns in the image.
+        
+    Returns:
+        list: List of (row, col) tuples.
+    """
+    return [(np.random.randint(ymin, ymax), np.random.randint(xmin, xmax)) for _ in range(n_pixels)]
+
+
+def random_pixel(xmin=0, xmax=512, ymin=0, ymax=1024):
+    """Return a random pixel.
+    
+    Returns:
+        tuple: (row, col)
+    """
+    return random_pixels(1, xmin, xmax, ymin, ymax)[0]

python/src/cluster.hpp

@@ -21,9 +21,9 @@ void define_cluster_finder_bindings(py::module &m) {
              })
         .def("pedestal",
              [](ClusterFinder<uint16_t, double> &self) {
-                 auto m = new NDArray<double, 2>{};
-                 *m = self.pedestal();
-                 return return_image_data(m);
+                 auto pd = new NDArray<double, 2>{};
+                 *pd = self.pedestal();
+                 return return_image_data(pd);
              })
         .def("find_clusters_without_threshold",
              [](ClusterFinder<uint16_t, double> &self,

python/src/ctb_raw_file.hpp

@@ -0,0 +1,57 @@
+
+#include "aare/CtbRawFile.hpp"
+#include "aare/File.hpp"
+#include "aare/Frame.hpp"
+#include "aare/RawFile.hpp"
+#include "aare/RawMasterFile.hpp"
+#include "aare/RawSubFile.hpp"
+
+#include "aare/defs.hpp"
+// #include "aare/fClusterFileV2.hpp"
+
+#include <cstdint>
+#include <filesystem>
+#include <pybind11/iostream.h>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl/filesystem.h>
+#include <string>
+
+namespace py = pybind11;
+using namespace ::aare;
+
+void define_ctb_raw_file_io_bindings(py::module &m) {
+
+ py::class_<CtbRawFile>(m, "CtbRawFile")
+        .def(py::init<const std::filesystem::path &>())
+        .def("read_frame",
+             [](CtbRawFile &self) {
+                 size_t image_size = self.image_size_in_bytes();
+                 py::array image;
+                 std::vector<ssize_t> shape;
+                 shape.reserve(2);
+                 shape.push_back(1);
+                 shape.push_back(image_size);
+
+                 py::array_t<DetectorHeader> header(1);
+
+                 // always read bytes
+                 image = py::array_t<uint8_t>(shape);
+
+                 self.read_into(
+                     reinterpret_cast<std::byte *>(image.mutable_data()),
+                     header.mutable_data());
+
+                 return py::make_tuple(header, image);
+             })
+        .def("seek", &CtbRawFile::seek)
+        .def("tell", &CtbRawFile::tell)
+        .def("master", &CtbRawFile::master)
+
+        .def_property_readonly("image_size_in_bytes",
+                               &CtbRawFile::image_size_in_bytes)
+
+        .def_property_readonly("frames_in_file", &CtbRawFile::frames_in_file);
+
+}

python/src/file.hpp

@@ -38,36 +38,7 @@ void define_file_io_bindings(py::module &m) {
                          bunchId, timestamp, modId, row, column, reserved,
                          debug, roundRNumber, detType, version, packetMask);
 
-    py::class_<CtbRawFile>(m, "CtbRawFile")
-        .def(py::init<const std::filesystem::path &>())
-        .def("read_frame",
-             [](CtbRawFile &self) {
-                 size_t image_size = self.image_size_in_bytes();
-                 py::array image;
-                 std::vector<ssize_t> shape;
-                 shape.reserve(2);
-                 shape.push_back(1);
-                 shape.push_back(image_size);
-
-                 py::array_t<DetectorHeader> header(1);
-
-                 // always read bytes
-                 image = py::array_t<uint8_t>(shape);
-
-                 self.read_into(
-                     reinterpret_cast<std::byte *>(image.mutable_data()),
-                     header.mutable_data());
-
-                 return py::make_tuple(header, image);
-             })
-        .def("seek", &CtbRawFile::seek)
-        .def("tell", &CtbRawFile::tell)
-        .def("master", &CtbRawFile::master)
-
-        .def_property_readonly("image_size_in_bytes",
-                               &CtbRawFile::image_size_in_bytes)
-
-        .def_property_readonly("frames_in_file", &CtbRawFile::frames_in_file);
+
 
     py::class_<File>(m, "File")
         .def(py::init([](const std::filesystem::path &fname) {
@@ -133,13 +104,15 @@ void define_file_io_bindings(py::module &m) {
                  return image;
              })
         .def("read_n", [](File &self, size_t n_frames) {
-            const uint8_t item_size = self.bytes_per_pixel();
+            //adjust for actual frames left in the file
+            n_frames = std::min(n_frames, self.total_frames()-self.tell());
+            if(n_frames == 0){
+                throw std::runtime_error("No frames left in file");
+            }
+            std::vector<size_t> shape{n_frames, self.rows(), self.cols()};
+
             py::array image;
-            std::vector<ssize_t> shape;
-            shape.reserve(3);
-            shape.push_back(n_frames);
-            shape.push_back(self.rows());
-            shape.push_back(self.cols());
+            const uint8_t item_size = self.bytes_per_pixel();
             if (item_size == 1) {
                 image = py::array_t<uint8_t>(shape);
             } else if (item_size == 2) {

python/src/module.cpp

@@ -1,6 +1,7 @@
 //Files with bindings to the different classes
 #include "file.hpp"
 #include "raw_file.hpp"
+#include "ctb_raw_file.hpp"
 #include "raw_master_file.hpp"
 #include "var_cluster.hpp"
 #include "pixel_map.hpp"
@@ -17,6 +18,7 @@ namespace py = pybind11;
 PYBIND11_MODULE(_aare, m) {
     define_file_io_bindings(m);
     define_raw_file_io_bindings(m);
+    define_ctb_raw_file_io_bindings(m);
     define_raw_master_file_bindings(m);
     define_var_cluster_finder_bindings(m);
     define_pixel_map_bindings(m);

python/src/pedestal.hpp

@@ -15,19 +15,19 @@ template <typename SUM_TYPE> void define_pedestal_bindings(py::module &m, const
         .def(py::init<int, int>())
         .def("mean",
              [](Pedestal<SUM_TYPE> &self) {
-                 auto m = new NDArray<SUM_TYPE, 2>{};
-                 *m = self.mean();
-                 return return_image_data(m);
+                 auto mea = new NDArray<SUM_TYPE, 2>{};
+                 *mea = self.mean();
+                 return return_image_data(mea);
              })
         .def("variance", [](Pedestal<SUM_TYPE> &self) {
-            auto m = new NDArray<SUM_TYPE, 2>{};
-            *m = self.variance();
-            return return_image_data(m);
+            auto var = new NDArray<SUM_TYPE, 2>{};
+            *var = self.variance();
+            return return_image_data(var);
         })
         .def("std", [](Pedestal<SUM_TYPE> &self) {
-            auto m = new NDArray<SUM_TYPE, 2>{};
-            *m = self.std();
-            return return_image_data(m);
+            auto std = new NDArray<SUM_TYPE, 2>{};
+            *std = self.std();
+            return return_image_data(std);
         })
         .def("clear", py::overload_cast<>(&Pedestal<SUM_TYPE>::clear))
         .def_property_readonly("rows", &Pedestal<SUM_TYPE>::rows)

python/src/raw_file.hpp

@@ -48,32 +48,44 @@ void define_raw_file_io_bindings(py::module &m) {
 
                  return py::make_tuple(header, image);
              })
-        .def("read_n",
-             [](RawFile &self, size_t n_frames) {
-                 py::array image;
-                 std::vector<ssize_t> shape;
-                 shape.reserve(3);
-                 shape.push_back(n_frames);
-                 shape.push_back(self.rows());
-                 shape.push_back(self.cols());
+        .def(
+            "read_n",
+            [](RawFile &self, size_t n_frames) {
+                // adjust for actual frames left in the file
+                n_frames =
+                    std::min(n_frames, self.total_frames() - self.tell());
+                if (n_frames == 0) {
+                    throw std::runtime_error("No frames left in file");
+                }
+                std::vector<size_t> shape{n_frames, self.rows(), self.cols()};
+
+                // return headers from all subfiles
+                py::array_t<DetectorHeader> header;
+                if (self.n_mod() == 1) {
+                    header = py::array_t<DetectorHeader>(n_frames);
+                } else {
+                    header = py::array_t<DetectorHeader>({self.n_mod(), n_frames});
+                }
+                // py::array_t<DetectorHeader> header({self.n_mod(), n_frames});
 
-                 // return headers from all subfiles
-                 py::array_t<DetectorHeader> header({self.n_mod(), n_frames});
-
-                 const uint8_t item_size = self.bytes_per_pixel();
-                 if (item_size == 1) {
-                     image = py::array_t<uint8_t>(shape);
-                 } else if (item_size == 2) {
-                     image = py::array_t<uint16_t>(shape);
-                 } else if (item_size == 4) {
-                     image = py::array_t<uint32_t>(shape);
-                 }
-                 self.read_into(
-                     reinterpret_cast<std::byte *>(image.mutable_data()),
-                     n_frames, header.mutable_data());
+                py::array image;
+                const uint8_t item_size = self.bytes_per_pixel();
+                if (item_size == 1) {
+                    image = py::array_t<uint8_t>(shape);
+                } else if (item_size == 2) {
+                    image = py::array_t<uint16_t>(shape);
+                } else if (item_size == 4) {
+                    image = py::array_t<uint32_t>(shape);
+                }
+                self.read_into(
+                    reinterpret_cast<std::byte *>(image.mutable_data()),
+                    n_frames, header.mutable_data());
 
-                 return py::make_tuple(header, image);
-             })
+                return py::make_tuple(header, image);
+            },
+            R"(
+             Read n frames from the file.
+             )")
         .def("frame_number", &RawFile::frame_number)
         .def_property_readonly("bytes_per_frame", &RawFile::bytes_per_frame)
         .def_property_readonly("pixels_per_frame", &RawFile::pixels_per_frame)