This commit is contained in:
Erik Fröjdh 2024-11-14 17:03:16 +01:00
parent 5cde7a99b5
commit 0d058274d5
2 changed files with 13 additions and 3 deletions

View File

@ -18,7 +18,7 @@ void define_cluster_file_io_bindings(py::module &m) {
PYBIND11_NUMPY_DTYPE(Cluster, x, y, data); PYBIND11_NUMPY_DTYPE(Cluster, x, y, data);
py::class_<ClusterFile>(m, "ClusterFile") py::class_<ClusterFile>(m, "ClusterFile")
.def(py::init<const std::filesystem::path &>()) .def(py::init<const std::filesystem::path &, size_t>(), py::arg(), py::arg("chunk_size") = 1000)
.def("read_clusters", .def("read_clusters",
[](ClusterFile &self, size_t n_clusters) { [](ClusterFile &self, size_t n_clusters) {
auto* vec = new std::vector<Cluster>(self.read_clusters(n_clusters)); auto* vec = new std::vector<Cluster>(self.read_clusters(n_clusters));
@ -29,6 +29,16 @@ void define_cluster_file_io_bindings(py::module &m) {
auto view = make_view_2d(noise_map); auto view = make_view_2d(noise_map);
auto* vec = new std::vector<Cluster>(self.read_cluster_with_cut(n_clusters, view.data(), nx, ny)); auto* vec = new std::vector<Cluster>(self.read_cluster_with_cut(n_clusters, view.data(), nx, ny));
return return_vector(vec); return return_vector(vec);
})
.def("__enter__", [](ClusterFile &self) { return &self; })
.def("__exit__", [](ClusterFile &self, py::args args) { return; })
.def("__iter__", [](ClusterFile &self) { return &self; })
.def("__next__", [](ClusterFile &self) {
auto vec = new std::vector<Cluster>(self.read_clusters(self.chunk_size()));
if(vec->size() == 0) {
throw py::stop_iteration();
}
return return_vector(vec);
}); });
} }

View File

@ -2,7 +2,7 @@
namespace aare { namespace aare {
ClusterFile::ClusterFile(const std::filesystem::path &fname) { ClusterFile::ClusterFile(const std::filesystem::path &fname, size_t chunk_size): m_chunk_size(chunk_size) {
fp = fopen(fname.c_str(), "rb"); fp = fopen(fname.c_str(), "rb");
if (!fp) { if (!fp) {
throw std::runtime_error("Could not open file: " + fname.string()); throw std::runtime_error("Could not open file: " + fname.string());