Adding array setting

henryiii · henryiii · commit 7b9b22a3a878 · 2019-11-25T16:08:06.000-05:00
diff --git a/README.md b/README.md
@@ -17,8 +17,6 @@ Python bindings for [Boost::Histogram][] ([source][Boost::Histogram source]), a
 > Please feel free to try out boost-histogram and give feedback.
 > Join the [discussion on gitter][gitter-link] or [open an issue](https://github.com/scikit-hep/boost-histogram/issues)!
 >
-> #### Known issues (develop):
-> * Setting with an array is not yet supported (`h[...] = np.array(...)`).
 
 
 ## Installation
diff --git a/boost_histogram/__init__.py b/boost_histogram/__init__.py
@@ -25,9 +25,8 @@
 try:
     from . import _core
 except ImportError as err:
-    err.msg += (
-        "\nDid you forget to compile? Use CMake or Setuptools to build, see the readme"
-    )
+    if "_core" in err.msg and "boost_histogram" in err.msg:
+        err.msg += "\nDid you forget to compile? Use CMake or Setuptools to build, see the readme"
     raise err
 
 
diff --git a/boost_histogram/_internal/hist.py b/boost_histogram/_internal/hist.py
@@ -465,8 +465,86 @@ def __getitem__(self, index):
             )
 
     def __setitem__(self, index, value):
-        indexes = self._compute_commonindex(index, expand_ellipsis=False)
-        self._hist._at_set(value, *indexes)
+        """
+        There are several supported possibilities:
+
+            h[slice] = array # same size
+
+        If an array is given to a compatible slice, it is set.
+
+            h[a:] = array # One larger
+
+        If an array is given that does not match, if it does match the
+        with-overflow size, it fills that.
+
+        PLANNED (not yet supported):
+
+            h[a:] = h2
+
+        If another histogram is given, that must either match with or without
+        overflow, where the overflow bins must be overflow bins (that is,
+        you cannot set a histogram's flow bins from another histogram that
+        is 2 larger). Bin edges must be a close match, as well. If you don't
+        want this level of type safety, just use ``h[...] = h2.view()``.
+        """
+        indexes = self._compute_commonindex(index, expand_ellipsis=True)
+
+        if isinstance(value, BaseHistogram):
+            raise TypeError("Not supported yet")
+
+        value = np.asarray(value)
+        view = self.view(flow=True)
+
+        # Disallow mismatched data types
+        if len(value.dtype) != len(view.dtype):
+            raise ValueError("Mismatched data types; matching types required")
+
+        # Numpy does not broadcast partial slices, but we would need
+        # to allow it (because we do allow broadcasting up dimensions)
+        # Instead, we simply require matching dimensions.
+        if value.ndim > 0 and value.ndim != len(indexes):
+            raise ValueError(
+                "Setting a histogram with an array must have a matching number of dimensions"
+            )
+
+        for n in range(len(indexes)):
+            request = indexes[n]
+            has_underflow = self.axes[n].options.underflow
+            has_overflow = self.axes[n].options.overflow
+
+            if isinstance(request, slice):
+                # Only consider underflow/overflow if the endpoints are not given
+                use_underflow = has_underflow and request.start is None
+                use_overflow = has_overflow and request.stop is None
+
+                # Make the limits explicit since we may need to shift them
+                start = 0 if request.start is None else request.start
+                stop = len(self.axes[n]) if request.stop is None else request.stop
+                request_len = stop - start
+
+                # If there are not enough dimensions, then treat it like broadcasting
+                if value.ndim == 0 or value.shape[n] == 1:
+                    start = 0 + has_overflow
+                    stop = len(self.axes[n]) + has_underflow
+                elif request_len == value.shape[n]:
+                    start += has_underflow
+                    stop += has_underflow
+                elif request_len + use_underflow + use_overflow == value.shape[n]:
+                    start += has_underflow and not use_underflow
+                    stop += has_underflow + (has_overflow and use_overflow)
+                else:
+                    msg = "Mismatched shapes in dimension {0}".format(n)
+                    msg += ", {0} != {1}".format(value.shape[n], request_len)
+                    if use_underflow or use_overflow:
+                        msg += " or {0}".format(
+                            request_len + use_underflow + use_overflow
+                        )
+                    raise ValueError(msg)
+                indexes[n] = slice(start, stop, request.step)
+            else:
+                indexes[n] = request + has_underflow
+
+        view[tuple(indexes)] = value
 
     def reduce(self, *args):
         """
diff --git a/include/histogram.hpp b/include/histogram.hpp
@@ -119,62 +119,3 @@ const F& pyarray_at(const py::array_t<F, Opt>& input,
     }
     return *(input.data() + flat_offset);
 }
-
-template <class Histogram>
-void copy_in(Histogram&, const py::array_t<double>&) {}
-
-template <>
-inline void copy_in<>(bh::histogram<vector_axis_variant, bh::dense_storage<double>>& h,
-                      const py::array_t<double>& input) {
-    // Works on simple datatypes only
-    // TODO: Add other types
-    if(h.rank() != input.ndim())
-        throw py::value_error(
-            "The input array dimensions must match the histogram rank");
-
-    // Quick check to ensure the input array is valid
-    for(unsigned r = 0; r < h.rank(); r++) {
-        auto input_shape = input.shape(static_cast<py::ssize_t>(r));
-        if(input_shape != bh::axis::traits::extent(h.axis(r))
-           && input_shape != h.axis(r).size() && input_shape != 1)
-            throw py::value_error("The input array sizes must match the histogram "
-                                  "(with or without flow), or be broadcastable to it");
-    }
-
-    std::vector<py::ssize_t> indexes;
-    indexes.resize(h.rank());
-
-    for(auto&& ind : bh::indexed(h, bh::coverage::all)) {
-        bool skip = false;
-
-        for(unsigned r = 0; r < h.rank(); r++) {
-            auto input_shape   = input.shape(static_cast<py::ssize_t>(r));
-            bool use_flow      = input_shape == bh::axis::traits::extent(h.axis(r));
-            bool has_underflow = h.axis(r).options() & bh::axis::option::underflow;
-
-            // Broadcast size 1
-            if(input_shape == 1)
-                indexes[r] = 0;
-
-            // If this is using flow bins and has an underflow bin, convert -1 to 0
-            // (etc)
-            else if(use_flow && has_underflow)
-                indexes[r] = ind.index(r) + 1;
-
-            // If not using flow bins, skip the flow bins
-            else if(!use_flow
-                    && (ind.index(r) < 0 || ind.index(r) >= h.axis(r).size())) {
-                skip = true;
-                break;
-
-                // Otherwise, this is normal
-            } else
-                indexes[r] = ind.index(r);
-        }
-
-        if(skip)
-            continue;
-
-        *ind = pyarray_at(input, indexes);
-    }
-}
diff --git a/include/register_histogram.hpp b/include/register_histogram.hpp
@@ -108,9 +108,6 @@ auto register_histogram(py::module& m, const char* name, const char* desc) {
             },
             "flow"_a = false)
 
-        .def("_copy_in",
-             [](histogram_t& h, py::array_t<double> input) { copy_in(h, input); })
-
         .def(
             "view",
             [](py::object self, bool flow) {
diff --git a/src/register_accumulators.cpp b/src/register_accumulators.cpp
@@ -51,6 +51,21 @@ decltype(auto) make_mean_call() {
     };
 }
 
+template <class T>
+decltype(auto) make_buffer() {
+    return [](T& self) -> py::buffer_info {
+        return py::buffer_info(
+            &self,                              // Pointer to buffer
+            sizeof(T),                          // Size of one scalar
+            py::format_descriptor<T>::format(), // Format registered with
+                                                // PYBIND11_NUMPY_DTYPE
+            0,                                  // Number of dimensions
+            {},                                 // Buffer dimensions
+            {}                                  // Stride
+        );
+    };
+}
+
 void register_accumulators(py::module& accumulators) {
     // Naming convention:
     // If a value is publically available in Boost.Histogram accumulators
@@ -62,7 +77,10 @@ void register_accumulators(py::module& accumulators) {
 
     PYBIND11_NUMPY_DTYPE(weighted_sum, value, variance);
 
-    register_accumulator<weighted_sum>(accumulators, "WeightedSum")
+    register_accumulator<weighted_sum>(
+        accumulators, "WeightedSum", py::buffer_protocol())
+
+        .def_buffer(make_buffer<weighted_sum>())
 
         .def(py::init<const double&>(), "value"_a)
         .def(py::init<const double&, const double&>(), "value"_a, "variance"_a)
@@ -158,7 +176,11 @@ void register_accumulators(py::module& accumulators) {
                          value,
                          sum_of_weighted_deltas_squared);
 
-    register_accumulator<weighted_mean>(accumulators, "WeightedMean")
+    register_accumulator<weighted_mean>(
+        accumulators, "WeightedMean", py::buffer_protocol())
+
+        .def_buffer(make_buffer<weighted_mean>())
+
         .def(py::init<const double&, const double&, const double&, const double&>(),
              "sum_of_weights"_a,
              "sum_of_weights_squared"_a,
@@ -238,7 +260,9 @@ void register_accumulators(py::module& accumulators) {
     using mean = accumulators::mean<double>;
     PYBIND11_NUMPY_DTYPE(mean, count, value, sum_of_deltas_squared);
 
-    register_accumulator<mean>(accumulators, "Mean")
+    register_accumulator<mean>(accumulators, "Mean", py::buffer_protocol())
+        .def_buffer(make_buffer<mean>())
+
         .def(py::init<const double&, const double&, const double&>(),
              "count"_a,
              "value"_a,
diff --git a/tests/test_histogram_set.py b/tests/test_histogram_set.py
@@ -0,0 +1,128 @@
+import boost_histogram as bh
+import numpy as np
+
+from numpy.testing import assert_array_equal
+
+import pytest
+
+
+def test_1D_set_bin():
+
+    h = bh.Histogram(bh.axis.Regular(10, 0, 1))
+
+    h[2] = 2
+    assert h[2] == 2.0
+
+    h[bh.underflow] = 3
+    assert h[bh.underflow] == 3.0
+
+    h[bh.overflow] = 4
+    assert h[bh.overflow] == 4.0
+
+
+def test_2d_set_bin():
+
+    h = bh.Histogram(bh.axis.Regular(10, 0, 1), bh.axis.Regular(10, 0, 1))
+
+    h[2, 4] = 2
+    assert h[2, 4] == 2.0
+
+    h[bh.underflow, 5] = 3
+    assert h[bh.underflow, 5] == 3.0
+
+    h[bh.overflow, bh.overflow] = 4
+    assert h[bh.overflow, bh.overflow] == 4.0
+
+
+def test_1d_set_array():
+    h = bh.Histogram(bh.axis.Regular(10, 0, 1))
+
+    h[...] = np.arange(10)
+    assert_array_equal(h.view(), np.arange(10))
+
+    h[...] = np.arange(12)
+    assert_array_equal(h.view(flow=True), np.arange(12))
+
+    with pytest.raises(ValueError):
+        h[...] = np.arange(9)
+    with pytest.raises(ValueError):
+        h[...] = np.arange(11)
+    with pytest.raises(ValueError):
+        h[...] = np.arange(13)
+
+    h[...] = 1
+    assert_array_equal(h.view(), np.ones(10))
+
+
+def test_2d_set_array():
+    h = bh.Histogram(bh.axis.Regular(10, 0, 1), bh.axis.Regular(10, 0, 1))
+
+    h[...] = np.arange(10).reshape(-1, 1)
+    assert_array_equal(h.view()[:, 2], np.arange(10))
+
+    h[...] = np.arange(12).reshape(-1, 1)
+    assert_array_equal(h.view(flow=True)[:, 3], np.arange(12))
+
+    with pytest.raises(ValueError):
+        h[...] = np.arange(9).reshape(-1, 1)
+    with pytest.raises(ValueError):
+        h[...] = np.arange(11).reshape(-1, 1)
+    with pytest.raises(ValueError):
+        h[...] = np.arange(13).reshape(-1, 1)
+
+    h[...] = 1
+    assert_array_equal(h.view(), np.ones((10, 10)))
+
+
+@pytest.mark.parametrize(
+    "storage, default",
+    (
+        (bh.storage.Mean, bh.accumulators.Mean(1.0, 2.0, 3.0)),
+        (bh.storage.WeightedMean, bh.accumulators.WeightedMean(1.0, 2.0, 3.0, 4.0)),
+        (bh.storage.Weight, bh.accumulators.WeightedSum(1.0, 2)),
+    ),
+)
+def test_set_special_dtype(storage, default):
+    h = bh.Histogram(
+        bh.axis.Regular(10, 0, 1), bh.axis.Regular(10, 0, 1), storage=storage()
+    )
+
+    arr = np.full((10, 1), default)
+    h[...] = arr
+    assert_array_equal(h.view()[:, 1:2], arr)
+
+    arr = np.full((12, 1), default)
+    h[...] = arr
+    assert_array_equal(h.view(flow=True)[:, 2:3], arr)
+
+    arr = np.full((10, 10), default)
+    h[...] = arr
+    assert_array_equal(h.view(), arr)
+
+    arr = np.full((10, 12), default)
+    h[...] = arr
+    assert_array_equal(h.view(flow=True)[1:11, :], arr)
+
+    arr = np.full((12, 10), default)
+    h[...] = arr
+    assert_array_equal(h.view(flow=True)[:, 1:11], arr)
+
+    arr = np.full((12, 12), default)
+    h[...] = arr
+    assert_array_equal(h.view(flow=True), arr)
+
+    with pytest.raises(ValueError):
+        arr = np.full((9, 1), default)
+        h[...] = arr
+    with pytest.raises(ValueError):
+        arr = np.full((11, 1), default)
+        h[...] = arr
+    with pytest.raises(ValueError):
+        arr = np.full((13, 1), default)
+        h[...] = arr
+
+    with pytest.raises(ValueError):
+        h[...] = 1
+
+    with pytest.raises(ValueError):
+        h[1, 1] = 1

Original file line number	Diff line number	Diff line change
`@@ -17,8 +17,6 @@ Python bindings for [Boost::Histogram][] ([source][Boost::Histogram source]), a`
`17`	`17`	`> Please feel free to try out boost-histogram and give feedback.`
`18`	`18`	`> Join the [discussion on gitter][gitter-link] or [open an issue](https://github.com/scikit-hep/boost-histogram/issues)!`
`19`	`19`	`>`
`20`		`-> #### Known issues (develop):`
`21`		-> * Setting with an array is not yet supported (`h[...] = np.array(...)`).
`22`	`20`
`23`	`21`
`24`	`22`	`## Installation`