diff --git a/pyproject.toml b/pyproject.toml index db42c57..186a1ab 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/quantity_array/__init__.py b/src/quantity_array/__init__.py index b4bc949..e846c2a 100644 --- a/src/quantity_array/__init__.py +++ b/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) diff --git a/tests/test_array.py b/tests/test_array.py index 6f3c478..92c8775 100644 --- a/tests/test_array.py +++ b/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])