feat: fix operators on arrays, adds support for it (#417)

* feat: support more forms of operators

* feat: shift work to the view as needed

* fix: better broadcasting, supports flow if present

Author: henryiii

include/bh_python/register_histogram.hpp

@@ -58,10 +58,7 @@ auto register_histogram(py::module& m, const char* name, const char* desc) {
                  return a;
              })
 
-        .def(py::self + py::self)
-        // .def(py::self + value_type())
         .def(py::self += py::self)
-        // .def(py::self += value_type())
 
         .def("__eq__",
              [](const histogram_t& self, const py::object& other) {
@@ -90,22 +87,21 @@ auto register_histogram(py::module& m, const char* name, const char* desc) {
 
         ;
 
-    // Atomics for example do not support these operations
+// Protection against an overzealous warning system
+// https://bugs.llvm.org/show_bug.cgi?id=43124
+#ifdef __clang__
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wself-assign-overloaded"
+#endif
     def_optionally(hist,
-                   bh::detail::has_operator_rmul<histogram_t, double>{},
-                   py::self *= double());
+                   bh::detail::has_operator_rdiv<histogram_t, histogram_t>{},
+                   py::self /= py::self);
     def_optionally(hist,
-                   bh::detail::has_operator_rmul<histogram_t, double>{},
-                   py::self * double());
-    def_optionally(hist,
-                   bh::detail::has_operator_rmul<histogram_t, double>{},
-                   double() * py::self);
-    def_optionally(hist,
-                   bh::detail::has_operator_rdiv<histogram_t, double>{},
-                   py::self /= double());
-    def_optionally(hist,
-                   bh::detail::has_operator_rdiv<histogram_t, double>{},
-                   py::self / double());
+                   bh::detail::has_operator_rmul<histogram_t, histogram_t>{},
+                   py::self *= py::self);
+#ifdef __clang__
+#pragma GCC diagnostic pop
+#endif
 
     hist.def(
             "to_numpy",

src/boost_histogram/_internal/hist.py

@@ -40,10 +40,6 @@ def _fill_cast(value):
         return np.ascontiguousarray(value)
 
 
-def _hist_or_val(other):
-    return other._hist if hasattr(other, "_hist") else other
-
-
 def _arg_shortcut(item):
     msg = "Developer shortcut: will be removed in a future version"
     if isinstance(item, tuple) and len(item) == 3:
@@ -185,22 +181,18 @@ def __array__(self):
         return self.view(False)
 
     def __add__(self, other):
-        if hasattr(other, "_hist"):
-            return self._new_hist(self._hist.__add__(other._hist))
-        else:
-            retval = self.copy()
-            retval += other
-            return retval
+        result = self.copy()
+        return result.__iadd__(other)
 
     def __iadd__(self, other):
         if isinstance(other, (int, float)) and other == 0:
             return self
-        if hasattr(other, "_hist"):
-            self._hist.__iadd__(other._hist)
-            # Addition may change category axes
-            self.axes = self._generate_axes_()
-            return self
-        return NotImplemented
+        self._compute_inplace_op("__iadd__", other)
+
+        # Addition may change the axes if they can grow
+        self.axes = self._generate_axes_()
+
+        return self
 
     def __radd__(self, other):
         return self + other
@@ -213,46 +205,55 @@ def __ne__(self, other):
 
     # If these fail, the underlying object throws the correct error
     def __mul__(self, other):
-        return self._new_hist(self._hist.__mul__(other))
+        result = self.copy()
+        return result._compute_inplace_op("__imul__", other)
 
     def __rmul__(self, other):
         return self * other
 
-    def __imul__(self, other):
-        self._hist.__imul__(_hist_or_val(other))
-        return self
-
     def __truediv__(self, other):
-        if isinstance(other, Histogram):
-            result = self.copy()
-            result.__itruediv__(other)
-            return result
-        else:
-            return self._new_hist(self._hist.__truediv__(_hist_or_val(other)))
+        result = self.copy()
+        return result._compute_inplace_op("__itruediv__", other)
 
     def __div__(self, other):
-        if isinstance(other, Histogram):
-            result = self.copy()
-            result.__idiv__(other)
-            return result
-        else:
-            return self._new_hist(self._hist.__div__(_hist_or_val(other)))
+        result = self.copy()
+        return result._compute_inplace_op("__idiv__", other)
 
-    def __itruediv__(self, other):
+    def _compute_inplace_op(self, name, other):
         if isinstance(other, Histogram):
-            view = self.view(flow=True)
-            view.__itruediv__(other.view(flow=True))
+            getattr(self._hist, name)(other._hist)
+        elif isinstance(other, _histograms):
+            getattr(self._hist, name)(other)
+        elif hasattr(other, "shape"):
+            if len(other.shape) != self.ndim:
+                raise ValueError(
+                    "Number of dimensions {0} must match histogram {1}".format(
+                        len(other.shape), self.ndim
+                    )
+                )
+            elif all((a == b or a == 1) for a, b in zip(other.shape, self.shape)):
+                view = self.view(flow=False)
+                getattr(view, name)(other)
+            elif all((a == b or a == 1) for a, b in zip(other.shape, self.axes.extent)):
+                view = self.view(flow=True)
+                getattr(view, name)(other)
+            else:
+                raise ValueError(
+                    "Wrong shape, expected {0} or {1}".format(self.shape, self.extent)
+                )
         else:
-            self._hist.__itruediv__(_hist_or_val(other))
+            view = self.view(flow=False)
+            getattr(view, name)(other)
         return self
 
     def __idiv__(self, other):
-        if isinstance(other, Histogram):
-            view = self.view(flow=True)
-            view.__idiv__(other.view(flow=True))
-        else:
-            self._hist.__idiv__(_hist_or_val(other))
-        return self
+        return self._compute_inplace_op("__idiv__", other)
+
+    def __itruediv__(self, other):
+        return self._compute_inplace_op("__itruediv__", other)
+
+    def __imul__(self, other):
+        return self._compute_inplace_op("__imul__", other)
 
     # TODO: Marked as too complex by flake8. Should be factored out a bit.
     @inject_signature("self, *args, weight=None, sample=None, threads=None")

tests/test_histogram.py

@@ -1078,3 +1078,86 @@ def test_shape():
 def test_empty_shape():
     h = bh.Histogram()
     assert h.shape == ()
+
+
+# issue #416 a
+def test_hist_division():
+    edges = [0, 0.25, 0.5, 0.75, 1, 2, 3, 4, 7, 10]
+    edges = [-x for x in reversed(edges)] + edges[1:]
+
+    h = bh.Histogram(bh.axis.Variable(edges))
+    h[...] = 1
+    h1 = h.copy()
+
+    dens = h.view().copy()
+    dens /= h.axes[0].widths * h.sum()
+
+    h1 /= h.axes[0].widths * h.sum()
+
+    assert_array_equal(h1.view(), dens)
+
+
+# issue #416 b
+# def test_hist_division():
+#     edges = [0, .25, .5, .75, 1, 2, 3, 4, 7, 10]
+#    edges = [-x for x in reversed(edges)] + edges[1:]
+#
+#    h = bh.Histogram(bh.axis.Variable(edges))
+#    h[...] = 1
+#
+#    dens = h.view().copy() / h.axes[0].widths * h.sum()
+#    h1 = h.copy()
+#
+#    h1[:] /=  h.axes[0].widths * h.sum()
+#
+#    assert_allclose(h1.view(), dens)
+
+
+def test_add_hists():
+    edges = [0, 0.25, 0.5, 0.75, 1, 2, 3, 4, 7, 10]
+    edges = [-x for x in reversed(edges)] + edges[1:]
+
+    h = bh.Histogram(bh.axis.Variable(edges))
+    h[...] = 1
+
+    h1 = h.copy()
+    h1 += h.view()
+
+    h2 = h.copy()
+    h2 += h1
+
+    h3 = h.copy()
+    h3 += 5
+
+    assert_array_equal(h, 1)
+    assert_array_equal(h1, 2)
+    assert_array_equal(h2, 3)
+    assert_array_equal(h3, 6)
+
+
+def test_add_broadcast():
+    h = bh.Histogram(bh.axis.Regular(10, 0, 1), bh.axis.Regular(20, 0, 1))
+
+    h1 = h.copy()
+    h2 = h.copy()
+
+    h1[...] = 1
+    assert h1.view().sum() == 10 * 20
+    assert h1.view(flow=True).sum() == 10 * 20
+
+    h2 = h + [[1]]
+    assert h2.sum() == 10 * 20
+    assert h2.sum(flow=True) == 10 * 20
+
+    h3 = h + np.ones((10, 20))
+    assert h3.sum() == 10 * 20
+    assert h3.sum(flow=True) == 10 * 20
+
+    h4 = h + np.ones((12, 22))
+    assert h4.view(flow=True).sum() == 12 * 22
+
+    h5 = h + np.ones((10, 1))
+    assert h5.sum(flow=True) == 10 * 20
+
+    h5 = h + np.ones((1, 22))
+    assert h5.sum(flow=True) == 12 * 22