diff --git a/docs/user-guide/storage.rst b/docs/user-guide/storage.rst index 908ef2b2..96721256 100644 --- a/docs/user-guide/storage.rst +++ b/docs/user-guide/storage.rst @@ -84,7 +84,7 @@ Mean This storage tracks a "Profile", that is, the mean value of the accumulation instead of the sum. It stores the count (as a double), the mean, and a term that is used to compute the variance. When -filling, you can add a ``sample=`` term. +filling, you must add a ``sample=`` term. WeightedMean diff --git a/include/bh_python/fill.hpp b/include/bh_python/fill.hpp index d2262f62..acdfb139 100644 --- a/include/bh_python/fill.hpp +++ b/include/bh_python/fill.hpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -95,6 +96,25 @@ inline decltype(auto) special_cast>(py::handle x) { return py::cast(x); } +// Allow single floats to be integers +template <> +inline decltype(auto) special_cast(py::handle x) { + try { + return static_cast(std::floor(py::cast(x))); + } catch(const py::cast_error&) { + return py::cast(x); + } +} + +// Allow float arrays to be integers +template <> +inline decltype(auto) special_cast>(py::handle x) { + auto np = py::module_::import("numpy"); + auto to_int = np.attr("floor")(x); + + return py::cast>(to_int); +} + using arg_t = variant::variant, double, c_array_t, diff --git a/tests/test_histogram.py b/tests/test_histogram.py index 30a608ff..607a7b95 100644 --- a/tests/test_histogram.py +++ b/tests/test_histogram.py @@ -117,6 +117,19 @@ def test_fill_int_1d(): h[-3] +def test_fill_int_with_float_single_1d(): + h = bh.Histogram(bh.axis.Integer(-1, 2)) + h.fill(0.3) + h.fill(-0.3) + assert h.values() == approx(np.array([1, 1, 0])) + + +def test_fill_int_with_float_array_1d(): + h = bh.Histogram(bh.axis.Integer(-1, 2)) + h.fill([-0.3, 0.3]) + assert h.values() == approx(np.array([1, 1, 0])) + + def test_fill_1d(flow): h = bh.Histogram(bh.axis.Regular(3, -1, 2, underflow=flow, overflow=flow)) with pytest.raises(ValueError):