From 8bf9ac55ced928dbfb181c88cf133a89ac9c77f9 Mon Sep 17 00:00:00 2001 From: froejdh_e Date: Wed, 27 Nov 2024 09:31:57 +0100 Subject: [PATCH] modified read_n also for File and RawFile --- python/CMakeLists.txt | 2 ++ python/aare/CtbRawFile.py | 9 ++--- python/aare/RawFile.py | 66 +++++++++++++++++++++++++++++++++++++ python/aare/__init__.py | 7 ++-- python/aare/utils.py | 23 +++++++++++++ python/src/cluster.hpp | 6 ++-- python/src/ctb_raw_file.hpp | 57 ++++++++++++++++++++++++++++++++ python/src/file.hpp | 45 +++++-------------------- python/src/module.cpp | 2 ++ python/src/pedestal.hpp | 18 +++++----- python/src/raw_file.hpp | 60 +++++++++++++++++++-------------- 11 files changed, 217 insertions(+), 78 deletions(-) create mode 100644 python/aare/RawFile.py create mode 100644 python/aare/utils.py create mode 100644 python/src/ctb_raw_file.hpp diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 888c33c..ab746dd 100644 --- a/python/CMakeLists.txt +++ b/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 diff --git a/python/aare/CtbRawFile.py b/python/aare/CtbRawFile.py index 00f4886..b2dcb0a 100644 --- a/python/aare/CtbRawFile.py +++ b/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 @@ class CtbRawFile(_aare.CtbRawFile): 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: diff --git a/python/aare/RawFile.py b/python/aare/RawFile.py new file mode 100644 index 0000000..87df440 --- /dev/null +++ b/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 diff --git a/python/aare/__init__.py b/python/aare/__init__.py index fced6b2..5641d85 100644 --- a/python/aare/__init__.py +++ b/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 \ No newline at end of file +from .RawFile import RawFile +from .ScanParameters import ScanParameters + +from .utils import random_pixels, random_pixel \ No newline at end of file diff --git a/python/aare/utils.py b/python/aare/utils.py new file mode 100644 index 0000000..d53f844 --- /dev/null +++ b/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] \ No newline at end of file diff --git a/python/src/cluster.hpp b/python/src/cluster.hpp index aa94b6a..6932281 100644 --- a/python/src/cluster.hpp +++ b/python/src/cluster.hpp @@ -21,9 +21,9 @@ void define_cluster_finder_bindings(py::module &m) { }) .def("pedestal", [](ClusterFinder &self) { - auto m = new NDArray{}; - *m = self.pedestal(); - return return_image_data(m); + auto pd = new NDArray{}; + *pd = self.pedestal(); + return return_image_data(pd); }) .def("find_clusters_without_threshold", [](ClusterFinder &self, diff --git a/python/src/ctb_raw_file.hpp b/python/src/ctb_raw_file.hpp new file mode 100644 index 0000000..39c1001 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; +using namespace ::aare; + +void define_ctb_raw_file_io_bindings(py::module &m) { + + py::class_(m, "CtbRawFile") + .def(py::init()) + .def("read_frame", + [](CtbRawFile &self) { + size_t image_size = self.image_size_in_bytes(); + py::array image; + std::vector shape; + shape.reserve(2); + shape.push_back(1); + shape.push_back(image_size); + + py::array_t header(1); + + // always read bytes + image = py::array_t(shape); + + self.read_into( + reinterpret_cast(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); + +} \ No newline at end of file diff --git a/python/src/file.hpp b/python/src/file.hpp index 5372618..3c44c43 100644 --- a/python/src/file.hpp +++ b/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_(m, "CtbRawFile") - .def(py::init()) - .def("read_frame", - [](CtbRawFile &self) { - size_t image_size = self.image_size_in_bytes(); - py::array image; - std::vector shape; - shape.reserve(2); - shape.push_back(1); - shape.push_back(image_size); - - py::array_t header(1); - - // always read bytes - image = py::array_t(shape); - - self.read_into( - reinterpret_cast(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_(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 shape{n_frames, self.rows(), self.cols()}; + py::array image; - std::vector 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(shape); } else if (item_size == 2) { diff --git a/python/src/module.cpp b/python/src/module.cpp index 1500425..7963ac4 100644 --- a/python/src/module.cpp +++ b/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); diff --git a/python/src/pedestal.hpp b/python/src/pedestal.hpp index 6cfcac3..4d5d043 100644 --- a/python/src/pedestal.hpp +++ b/python/src/pedestal.hpp @@ -15,19 +15,19 @@ template void define_pedestal_bindings(py::module &m, const .def(py::init()) .def("mean", [](Pedestal &self) { - auto m = new NDArray{}; - *m = self.mean(); - return return_image_data(m); + auto mea = new NDArray{}; + *mea = self.mean(); + return return_image_data(mea); }) .def("variance", [](Pedestal &self) { - auto m = new NDArray{}; - *m = self.variance(); - return return_image_data(m); + auto var = new NDArray{}; + *var = self.variance(); + return return_image_data(var); }) .def("std", [](Pedestal &self) { - auto m = new NDArray{}; - *m = self.std(); - return return_image_data(m); + auto std = new NDArray{}; + *std = self.std(); + return return_image_data(std); }) .def("clear", py::overload_cast<>(&Pedestal::clear)) .def_property_readonly("rows", &Pedestal::rows) diff --git a/python/src/raw_file.hpp b/python/src/raw_file.hpp index 73f7699..38b4896 100644 --- a/python/src/raw_file.hpp +++ b/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 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 shape{n_frames, self.rows(), self.cols()}; + + // return headers from all subfiles + py::array_t header; + if (self.n_mod() == 1) { + header = py::array_t(n_frames); + } else { + header = py::array_t({self.n_mod(), n_frames}); + } + // py::array_t header({self.n_mod(), n_frames}); - // return headers from all subfiles - py::array_t header({self.n_mod(), n_frames}); + py::array image; + const uint8_t item_size = self.bytes_per_pixel(); + if (item_size == 1) { + image = py::array_t(shape); + } else if (item_size == 2) { + image = py::array_t(shape); + } else if (item_size == 4) { + image = py::array_t(shape); + } + self.read_into( + reinterpret_cast(image.mutable_data()), + n_frames, header.mutable_data()); - const uint8_t item_size = self.bytes_per_pixel(); - if (item_size == 1) { - image = py::array_t(shape); - } else if (item_size == 2) { - image = py::array_t(shape); - } else if (item_size == 4) { - image = py::array_t(shape); - } - self.read_into( - reinterpret_cast(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)