Commit 2cc167b3 authored by Roel Aaij's avatar Roel Aaij
Browse files

Update output container shape to be (4, n)

parent 07fc1a1b
......@@ -17,7 +17,11 @@
#include <generate.h>
#include <storage.h>
namespace py = pybind11;
namespace {
namespace py = pybind11;
using std::vector;
using std::array;
}
inline xt::pyarray<long>
generate_k40(const long time_start, const long time_end, Generators& gens, bool avx2)
......@@ -26,9 +30,35 @@ generate_k40(const long time_start, const long time_end, Generators& gens, bool
auto times = std::get<0>(r);
auto values = std::get<1>(r);
auto n = times.size();
times.resize(2 * n);
std::copy(begin(values), end(values), begin(times) + n);
return xt::adapt(std::move(times), std::array<size_t, 2>{2, n});
// Create output array and copy timestamps into it
vector<long> output(4 * n);
std::copy(begin(times), end(times), begin(output));
// factory for lambda's to shift and mask the values back into their
// components.
auto shift_mask_fact = [&values] (const size_t shift, const long mask) {
return [&values, shift, mask] (const auto val) {
return (val >> shift) & mask;
};
};
// column in the output array
size_t column = 1;
// Add DOM ID, PMT ID and TOT to the output array
for (auto fun : {// DOM ID
shift_mask_fact(13, 0xffff),
// PMT ID
shift_mask_fact(8, 0x1f),
// TOT
shift_mask_fact(0, 0xff)}) {
for (size_t i = 0; i < values.size(); ++i) {
output[values.size() * column + i] = fun(values[i]);
}
++column;
}
// Adapt to an xt::array to give to python
return xt::adapt(std::move(output), array<size_t, 2>{4, n});
}
// Python Module and Docstrings
......@@ -37,7 +67,7 @@ PYBIND11_MODULE(k40gen, m)
xt::import_numpy();
py::class_<Generators>(m, "Generators")
.def(py::init<const uint, const uint, std::array<float, 4>>());
.def(py::init<const uint, const uint, array<float, 4>>());
m.doc() = R"pbdoc(
......@@ -54,5 +84,6 @@ PYBIND11_MODULE(k40gen, m)
m.def("generate_k40", &generate_k40,
py::call_guard<py::scoped_ostream_redirect,
py::scoped_estream_redirect>{},
"Generate K40 background");
"Generate K40 background as a (4, n) array with columns: "
"(time, DOM ID, PMT ID, time-over-threshold)");
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment