Repository: hgrecco/pint

Issue:
Quantity.prod() result unit is incorrect with axis or where argument
Right now, the `.prod()` method assumes the full size of the input is collapsed in the result. This gives incorrect results when the `axis` or `where` arguments are supplied, as seen below:

```python
import pint

ureg = pint.UnitRegistry()

q = [[1, 2], [3, 4]] * ureg.m

print(q.prod())
print(q.prod(axis=0))
print(q.prod(where=[True, False]))
```
```
24 meter ** 4
[3 8] meter ** 4
3 meter ** 4
```

The unit on the first result where the full array is collapsed to a scalar is correct, but the other two results should have `meter ** 2` as the unit.

I don't have a good fix for this yet, since this is the same problem mentioned in https://github.com/hgrecco/pint/pull/764#issuecomment-523749008. If anyone has any suggestions for a performant way to determine how many unit multiplications occur given both `axis` and `where` arguments, please do let me know! Otherwise, I'll try to figure something out and get a PR in for this or include it alongside a `prod` implementation for `__array_function__`.

Discussion:

Patch:
diff --git a/pint/numpy_func.py b/pint/numpy_func.py
index 5994461..c671338 100644
--- a/pint/numpy_func.py
+++ b/pint/numpy_func.py
@@ -677,21 +677,27 @@ def _prod(a, *args, **kwargs):
     axis = all_kwargs.get("axis", None)
     where = all_kwargs.get("where", None)
 
-    if axis is not None and where is not None:
-        raise ValueError("passing axis and where is not supported")
-
-    result = np.prod(a._magnitude, *args, **kwargs)
+    registry = a.units._REGISTRY
 
-    if axis is not None:
-        exponent = a.size // result.size
-        units = a.units ** exponent
+    if axis is not None and where is not None:
+        _, where_ = np.broadcast_arrays(a._magnitude, where)
+        exponents = np.unique(np.sum(where_, axis=axis))
+        if len(exponents) == 1 or (len(exponents) == 2 and 0 in exponents):
+            units = a.units ** np.max(exponents)
+        else:
+            units = registry.dimensionless
+            a = a.to(units)
+    elif axis is not None:
+        units = a.units ** a.shape[axis]
     elif where is not None:
-        exponent = np.asarray(where, dtype=np.bool_).sum()
+        exponent = np.sum(where)
         units = a.units ** exponent
     else:
         units = a.units ** a.size
 
-    return units._REGISTRY.Quantity(result, units)
+    result = np.prod(a._magnitude, *args, **kwargs)
+
+    return registry.Quantity(result, units)
 
 
 # Implement simple matching-unit or stripped-unit functions based on signature


Test patch:
diff --git a/pint/testsuite/test_numpy.py b/pint/testsuite/test_numpy.py
index 9ecdcb1..ea570ef 100644
--- a/pint/testsuite/test_numpy.py
+++ b/pint/testsuite/test_numpy.py
@@ -295,7 +295,14 @@ class TestNumpyMathematicalFunctions(TestNumpyMethods):
         self.assertQuantityEqual(np.prod(self.q, axis=axis), [3, 8] * self.ureg.m ** 2)
         self.assertQuantityEqual(np.prod(self.q, where=where), 12 * self.ureg.m ** 3)
 
-        self.assertRaises(ValueError, np.prod, self.q, axis=axis, where=where)
+        self.assertRaises(DimensionalityError, np.prod, self.q, axis=axis, where=where)
+        self.assertQuantityEqual(
+            np.prod(self.q, axis=axis, where=[[True, False], [False, True]]),
+            [1, 4] * self.ureg.m,
+        )
+        self.assertQuantityEqual(
+            np.prod(self.q, axis=axis, where=[True, False]), [3, 1] * self.ureg.m ** 2
+        )
 
     def test_sum(self):
         self.assertEqual(self.q.sum(), 10 * self.ureg.m)