update counts() and variances()

Author: siliataider

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi.py

@@ -519,34 +519,34 @@ def _values_by_copy(self, flow=False) -> np.typing.NDArray[Any]:  # noqa: F821
     return np.array([self.GetBinContent(*bin) for bin in bin_combinations]).reshape(_shape(self, flow=flow))
 
 
-def _variances(self) -> np.typing.NDArray[Any]:  # noqa: F821
+def _variances(self, flow=False) -> np.typing.NDArray[Any]:  # noqa: F821
     import numpy as np
 
-    sum_of_weights = self.values()
+    sum_of_weights = self.values(flow=flow)
 
     if not _hasWeights(self) and _kind(self) == Kind.COUNT:
         return sum_of_weights
 
-    sum_of_weights_squared = _get_sum_of_weights_squared(self)
+    sum_of_weights_squared = _get_sum_of_weights_squared(self, flow=flow)
 
     if _kind(self) == Kind.MEAN:
-        counts = self.counts()
+        counts = self.counts(flow=flow)
         variances = sum_of_weights_squared.copy()
         variances[counts <= 1] = np.nan
         return variances
 
     return sum_of_weights_squared
 
 
-def _counts(self) -> np.typing.NDArray[Any]:  # noqa: F821
+def _counts(self, flow=False) -> np.typing.NDArray[Any]:  # noqa: F821
     import numpy as np
 
-    sum_of_weights = self.values()
+    sum_of_weights = self.values(flow=flow)
 
     if not _hasWeights(self):
         return sum_of_weights
 
-    sum_of_weights_squared = _get_sum_of_weights_squared(self)
+    sum_of_weights_squared = _get_sum_of_weights_squared(self, flow=flow)
 
     return np.divide(
         sum_of_weights**2,
@@ -556,16 +556,23 @@ def _counts(self) -> np.typing.NDArray[Any]:  # noqa: F821
     )
 
 
-def _get_sum_of_weights_squared(self) -> np.typing.NDArray[Any]:  # noqa: F821
+def _get_sum_of_weights_squared(self, flow=False) -> np.typing.NDArray[Any]:  # noqa: F821
     import numpy as np
 
-    shape = _shape(self, flow=False)
     sumw2_arr = np.frombuffer(
         self.GetSumw2().GetArray(),
         dtype=self.GetSumw2().GetArray().typecode,
         count=self.GetSumw2().GetSize(),
     )
-    return sumw2_arr[tuple([slice(1, -1)] * len(shape))].reshape(shape, order="F") if sumw2_arr.size > 0 else sumw2_arr
+
+    reshaped = sumw2_arr.reshape(_shape(self, flow=True), order="F")
+
+    if flow:
+        slices = tuple(slice(None) for _ in range(self.GetDimension()))
+    else:
+        slices = tuple(slice(1, -1) for _ in range(self.GetDimension()))
+
+    return reshaped[slices]
 
 
 values_func_dict: dict[str, Callable] = {