diff --git a/src/python/histogram.cpp b/src/python/histogram.cpp index 561c232f..8c0fa58d 100644 --- a/src/python/histogram.cpp +++ b/src/python/histogram.cpp @@ -342,12 +342,37 @@ python::object histogram_variance(python::tuple args, python::dict kwargs) { } int idx[BOOST_HISTOGRAM_AXIS_LIMIT]; - for (unsigned i = 0; i < self.dim(); ++i) + for (auto i = 0u; i < self.dim(); ++i) idx[i] = python::extract(args[1 + i]); return python::object(self.variance(idx, idx + self.dim())); } +python::object histogram_reduce_to(python::tuple args, python::dict kwargs) { + const histogram &self = + python::extract(args[0]); + + const unsigned nargs = len(args) - 1; + + if (nargs > BOOST_HISTOGRAM_AXIS_LIMIT) { + std::ostringstream os; + os << "too many arguments, maximum is " << BOOST_HISTOGRAM_AXIS_LIMIT; + PyErr_SetString(PyExc_RuntimeError, os.str().c_str()); + python::throw_error_already_set(); + } + + if (kwargs) { + PyErr_SetString(PyExc_RuntimeError, "no keyword arguments allowed"); + python::throw_error_already_set(); + } + + int idx[BOOST_HISTOGRAM_AXIS_LIMIT]; + for (auto i = 0u; i < nargs; ++i) + idx[i] = python::extract(args[1 + i]); + + return python::object(self.reduce_to(idx, idx + nargs)); +} + std::string histogram_repr(const histogram &h) { std::ostringstream os; os << h; @@ -389,6 +414,9 @@ void register_histogram() { .def("variance", python::raw_function(histogram_variance), ":param int args: indices of the bin (number must match dimension)" "\n:return: variance estimate for the bin") + .def("reduce_to", python::raw_function(histogram_reduce_to), + ":param int args: indices of the axes in the reduced histogram" + "\n:return: reduced histogram with subset of axes") .def("__repr__", histogram_repr, ":return: string representation of the histogram") .def(python::self == python::self) diff --git a/test/python_suite_test.py b/test/python_suite_test.py index 48a5407d..e05d8db3 100644 --- a/test/python_suite_test.py +++ b/test/python_suite_test.py @@ -603,6 +603,22 @@ class test_histogram(unittest.TestCase): with self.assertRaises(RuntimeError): h + h2 + def test_reduce_to(self): + h = histogram(integer(0, 2), integer(1, 4)) + h.fill(0, 1) + h.fill(0, 2) + h.fill(1, 3) + + h0 = h.reduce_to(0) + self.assertEqual(h0.dim, 1) + self.assertEqual(h0.axis(), integer(0, 2)) + self.assertEqual([h0.value(i) for i in range(2)], [2, 1]) + + h1 = h.reduce_to(1) + self.assertEqual(h1.dim, 1) + self.assertEqual(h1.axis(), integer(1, 4)) + self.assertEqual([h1.value(i) for i in range(3)], [1, 1, 1]) + def test_pickle_0(self): a = histogram(category(0, 1, 2), integer(0, 20, label='ia'), diff --git a/test/speed_cpp.cpp b/test/speed_cpp.cpp index 36289ef1..64f47a8c 100644 --- a/test/speed_cpp.cpp +++ b/test/speed_cpp.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include