Fix units returned by multiply, divide/floor_divide, and square

pyproject.toml

@@ -69,6 +69,10 @@ array-api-strict = "*"
 # clean array-api-tests dir
 clean-xp-tests = { cwd = ".", cmd = "rm -rf array-api-tests" }
 # clone array-api-tests
+# The checkout is pinned: the test suite tracks the array API standard and
+# regularly adds tests that the pinned array-api-strict does not implement yet
+# (finfo/iinfo with an array argument, for example).  Pinning keeps the suite
+# in step with array-api-strict so that a failure means a real regression.
 clone-xp-tests.cmd = [
     "git",
     "clone",
@@ -78,6 +82,10 @@ clone-xp-tests.cmd = [
     "array-api-tests",
     "&&",
     "git",
+    "checkout",
+    "c48410f96fc58e02eea844e6b7f6cc01680f77ce",
+    "&&",
+    "git",
     "submodule",
     "update",
     "--init",

src/quantity_array/__init__.py

@@ -675,7 +675,6 @@ def where(condition, x1, x2, /):
         "signbit",
         "sin",
         "sinh",
-        "square",
         "tan",
         "tanh",
         "trunc",
@@ -698,6 +697,14 @@ def _sqrt(x, /, *args, **kwargs):
 
     mod.sqrt = _sqrt
 
+    def _square(x, /, *args, **kwargs):
+        x = asarray(x)
+        magnitude = xp.asarray(x.magnitude, copy=True)
+        magnitude = xp.square(magnitude, *args, **kwargs)
+        return ArrayUnitQuantity(magnitude, x.units**2)
+
+    mod.square = _square
+
     elementwise_two_arrays = [
         "add",
         "atan2",
@@ -707,9 +714,7 @@ def _sqrt(x, /, *args, **kwargs):
         "bitwise_right_shift",
         "bitwise_xor",
         "copysign",
-        "divide",
         "equal",
-        "floor_divide",
         "greater",
         "greater_equal",
         "hypot",
@@ -749,13 +754,52 @@ def multiply(x1, x2, /, *args, **kwargs):
         units = x1.units * x2.units
 
         x1_magnitude = xp.asarray(x1.magnitude, copy=True)
-        x2_magnitude = x2.m_as(x1.units)
+        x2_magnitude = x2.magnitude
 
         magnitude = xp.multiply(x1_magnitude, x2_magnitude, *args, **kwargs)
         return ArrayUnitQuantity(magnitude, units)
 
     mod.multiply = multiply
 
+    def divide(x1, x2, /, *args, **kwargs):
+        x1 = asarray(x1)
+        x2 = asarray(x2)
+
+        units = x1.units / x2.units
+
+        x1_magnitude = xp.asarray(x1.magnitude, copy=True)
+        x2_magnitude = x2.magnitude
+
+        magnitude = xp.divide(x1_magnitude, x2_magnitude, *args, **kwargs)
+        return ArrayUnitQuantity(magnitude, units)
+
+    mod.divide = divide
+
+    def floor_divide(x1, x2, /, *args, **kwargs):
+        x1 = asarray(x1)
+        x2 = asarray(x2)
+
+        x1_magnitude = xp.asarray(x1.magnitude, copy=True)
+        x2_magnitude = xp.asarray(x2.magnitude, copy=True)
+
+        if x1.units != x2.units:
+            # The units differ, so the magnitudes are not commensurate and
+            # have to be brought onto a common scale before flooring.
+            # m_as() raises DimensionalityError for dimensionally
+            # incompatible units, as pint's Quantity.__floordiv__ does.
+            # The conversion factor is generally fractional, so this is done
+            # in floating point.
+            x1_magnitude = xp.astype(x1_magnitude, xp.float64)
+            x2_magnitude = astype(x2, xp.float64).m_as(x1.units)
+
+        magnitude = xp.floor_divide(x1_magnitude, x2_magnitude, *args, **kwargs)
+        # Floor division of two quantities is a count, not a ratio: it is the
+        # number of whole right-hand-side units that fit in one left-hand-side
+        # unit, so the result is dimensionless.
+        return ArrayUnitQuantity(magnitude, None)
+
+    mod.floor_divide = floor_divide
+
     def pow(x1, x2, /, *args, **kwargs):
         x1 = asarray(x1)
         x2 = asarray(x2)
@@ -931,7 +975,16 @@ def var(x, /, *args, **kwargs):
     def prod(x, /, *args, axis=None, **kwargs):
         x = asarray(x)
         magnitude = xp.asarray(x.magnitude, copy=True)
-        exponent = magnitude.shape[axis] if axis is not None else size(magnitude)
+        if axis is None:
+            exponent = size(magnitude)
+        elif isinstance(axis, tuple):
+            # The product reduces over every given axis, so the unit is
+            # raised to the power of the total number of elements reduced.
+            exponent = 1
+            for axis_i in axis:
+                exponent *= magnitude.shape[axis_i]
+        else:
+            exponent = magnitude.shape[axis]
         units = x.units**exponent
         magnitude = xp.prod(magnitude, *args, axis=axis, **kwargs)
         return ArrayUnitQuantity(magnitude, units)

tests/test_array.py

@@ -525,3 +525,112 @@ def test_tile(self):
             pxp.tile(pxp.asarray([[1, 2], [3, 4]]) * self.ureg.m, (2, 1)),
             pxp.asarray([[1, 2], [3, 4], [1, 2], [3, 4]]) * self.ureg.m,
         )
+
+
+class TestMultiplyUnits(TestArrayMethods):
+    def test_multiply_units(self):
+        result = pxp.multiply(
+            pxp.asarray([1.0], units="m"), pxp.asarray([1.0], units="mm")
+        )
+        assert result.units == self.ureg.m * self.ureg.mm
+        np.testing.assert_allclose(np.asarray(result.m_as("m**2"), dtype=float), [1e-3])
+
+        result = pxp.multiply(
+            pxp.asarray([2.0], units="m"), pxp.asarray([3.0], units="s")
+        )
+        assert result.units == self.ureg.m * self.ureg.s
+        np.testing.assert_allclose(np.asarray(result.m_as("m*s"), dtype=float), [6.0])
+
+
+class TestDivideUnits(TestArrayMethods):
+    def test_divide_units(self):
+        result = pxp.divide(
+            pxp.asarray([2.0], units="m"), pxp.asarray([1.0], units="m")
+        )
+        assert result.units == self.ureg.dimensionless
+        np.testing.assert_allclose(np.asarray(result.m_as(""), dtype=float), [2.0])
+
+        result = pxp.divide(
+            pxp.asarray([6.0], units="m"), pxp.asarray([3.0], units="s")
+        )
+        assert result.units == self.ureg.m / self.ureg.s
+        np.testing.assert_allclose(np.asarray(result.m_as("m/s"), dtype=float), [2.0])
+
+    def test_floor_divide_units(self):
+        result = pxp.floor_divide(
+            pxp.asarray([7.0], units="m"), pxp.asarray([2.0], units="m")
+        )
+        assert result.units == self.ureg.dimensionless
+        np.testing.assert_allclose(np.asarray(result.m_as(""), dtype=float), [3.0])
+
+    def test_floor_divide_scale_factor(self):
+        # The units differ by a scale factor, so the magnitudes have to be
+        # brought onto a common scale before flooring: 2 m is 200 cm, and
+        # 200 cm // 3 cm is 66.
+        result = pxp.floor_divide(
+            pxp.asarray([2.0], units="m"), pxp.asarray([3.0], units="cm")
+        )
+        assert result.units == self.ureg.dimensionless
+        np.testing.assert_allclose(np.asarray(result.m_as(""), dtype=float), [66.0])
+
+    def test_floor_divide_scale_factor_is_invariant(self):
+        # The same physical ratio written with different units must give the
+        # same answer.
+        results = [
+            pxp.floor_divide(
+                pxp.asarray([2.0], units="m"), pxp.asarray([3.0], units="cm")
+            ),
+            pxp.floor_divide(
+                pxp.asarray([200.0], units="cm"), pxp.asarray([3.0], units="cm")
+            ),
+            pxp.floor_divide(
+                pxp.asarray([2000.0], units="mm"), pxp.asarray([3.0], units="cm")
+            ),
+        ]
+        for result in results:
+            assert result.units == self.ureg.dimensionless
+            np.testing.assert_allclose(np.asarray(result.m_as(""), dtype=float), [66.0])
+
+    def test_floor_divide_matches_operator(self):
+        # The function and the operator must agree.
+        x1 = pxp.asarray([2.0], units="m")
+        x2 = pxp.asarray([3.0], units="cm")
+        np.testing.assert_allclose(
+            np.asarray(pxp.floor_divide(x1, x2).magnitude, dtype=float),
+            np.asarray(x1 // x2, dtype=float),
+        )
+
+    def test_floor_divide_incompatible_units(self):
+        # Dimensionally incompatible units cannot be floor divided.
+        with pytest.raises(DimensionalityError):
+            pxp.floor_divide(
+                pxp.asarray([6.0], units="m"), pxp.asarray([4.0], units="s")
+            )
+
+
+class TestSquareUnits(TestArrayMethods):
+    def test_square_units(self):
+        result = pxp.square(pxp.asarray([3.0], units="m"))
+        assert result.units == self.ureg.m**2
+        np.testing.assert_allclose(np.asarray(result.m_as("m**2"), dtype=float), [9.0])
+
+
+class TestProdUnits(TestArrayMethods):
+    def test_prod_units(self):
+        result = pxp.prod(pxp.asarray([[1.0, 2.0], [3.0, 4.0]], units="m"))
+        assert result.units == self.ureg.m**4
+        np.testing.assert_allclose(np.asarray(result.m_as("m**4"), dtype=float), [24.0])
+
+    def test_prod_units_axis(self):
+        result = pxp.prod(pxp.asarray([[1.0, 2.0], [3.0, 4.0]], units="m"), axis=0)
+        assert result.units == self.ureg.m**2
+        np.testing.assert_allclose(
+            np.asarray(result.m_as("m**2"), dtype=float), [3.0, 8.0]
+        )
+
+    def test_prod_units_tuple_axis(self):
+        # Reducing over several axes raises the unit to the power of the
+        # total number of elements reduced.
+        result = pxp.prod(pxp.asarray([[1.0, 2.0], [3.0, 4.0]], units="m"), axis=(0, 1))
+        assert result.units == self.ureg.m**4
+        np.testing.assert_allclose(np.asarray(result.m_as("m**4"), dtype=float), [24.0])