diff --git a/CHANGES.txt b/CHANGES.txt index 73e05ed..600922c 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -6,6 +6,9 @@ CHANGES Unreleased ----------- +- Support matrix multiplication with ``@`` and ``numpy.matmul``, multiplying + the operand units (`#231 `_). + - Correct the pond force unit to 0.00980665 N (one gram-force); it was previously 1000 times too large (`#44 `_). diff --git a/doc/user/tutorial.rst b/doc/user/tutorial.rst index e29d304..d6bf779 100644 --- a/doc/user/tutorial.rst +++ b/doc/user/tutorial.rst @@ -3,3 +3,20 @@ Quick-start tutorial ==================== .. automodule:: quantities + +Matrix multiplication +===================== + +The ``@`` operator and ``numpy.matmul`` multiply the units of their operands, +just like elementwise multiplication:: + + >>> import quantities as pq + >>> left = pq.Quantity([[1, 2], [3, 4]], 'm') + >>> right = pq.Quantity([[5, 6], [7, 8]], 's') + >>> result = left @ right + >>> result.magnitude.tolist() + [[19, 22], [43, 50]] + >>> bool(result.units == pq.m * pq.s) + True + +Uncertainty propagation for matrix multiplication is not supported. diff --git a/quantities/dimensionality.py b/quantities/dimensionality.py index f9cc3e2..d816762 100644 --- a/quantities/dimensionality.py +++ b/quantities/dimensionality.py @@ -227,6 +227,16 @@ def _d_multiply(q1, q2, out=None): p_dict[np.multiply] = _d_multiply p_dict[np.cross] = _d_multiply +def _d_matmul(q1, q2, out=None): + from .uncertainquantity import UncertainQuantity + + if any(isinstance(q, UncertainQuantity) for q in (q1, q2, out)): + raise ValueError('matrix multiplication is not supported for UncertainQuantity') + if not hasattr(q1, 'dimensionality') and not hasattr(q2, 'dimensionality'): + return Dimensionality() + return _d_multiply(q1, q2, out) +p_dict[np.matmul] = _d_matmul + def _d_divide(q1, q2, out=None): try: return q1._dimensionality / q2._dimensionality diff --git a/quantities/quantity.py b/quantities/quantity.py index 0cde19c..db68d82 100644 --- a/quantities/quantity.py +++ b/quantities/quantity.py @@ -324,7 +324,8 @@ def __array_wrap__(self, obj, context=None, return_scalar=False): # For NumPy < 2.0 we do old behavior if _np_version < (2, 0, 0): - if not isinstance(obj, Quantity): + # Generalized ufuncs do not call __array_prepare__, even for out. + if not isinstance(obj, Quantity) or (context is not None and context[0] is np.matmul): return self.__array_prepare__(obj, context) else: return obj diff --git a/quantities/quantity.pyi b/quantities/quantity.pyi index 2d5db95..526482e 100644 --- a/quantities/quantity.pyi +++ b/quantities/quantity.pyi @@ -70,7 +70,8 @@ class Quantity(npt.NDArray): def __rmul__(self, other) -> Quantity: ... def __imul__(self, other) -> Quantity: ... - # NOTE matmul is not supported + def __matmul__(self, other) -> Quantity: ... + def __rmatmul__(self, other) -> Quantity: ... def __truediv__(self, other) -> Quantity: ... # type: ignore[override] def __rtruediv__(self, other) -> Quantity: ... # type: ignore[override] diff --git a/quantities/tests/test_matmul.py b/quantities/tests/test_matmul.py new file mode 100644 index 0000000..cea5787 --- /dev/null +++ b/quantities/tests/test_matmul.py @@ -0,0 +1,108 @@ +import operator + +import numpy as np + +from .. import Quantity, UncertainQuantity +from .. import units as pq +from .common import TestCase, unittest + + +class TestMatmul(TestCase): + + def test_matrix_multiplication(self): + q = Quantity([[1, 2, 3], [4, 5, 6]], 'um') + expected = [[14, 32], [32, 77]] * pq.um**2 + self.assertQuantityEqual(q @ q.T, expected) + self.assertQuantityEqual(np.matmul(q, q.T), expected) + self.assertQuantityEqual(q, [[1, 2, 3], [4, 5, 6]] * pq.um) + + def test_vectors_and_broadcasting(self): + for left_shape, right_shape in ( + ((3,), (3,)), + ((2, 3), (3,)), + ((3,), (3, 4)), + ((2, 3), (3, 4)), + ((2, 1, 3, 4), (5, 4, 2))): + with self.subTest(left_shape=left_shape, right_shape=right_shape): + left = np.arange(np.prod(left_shape)).reshape(left_shape) + right = np.arange(np.prod(right_shape)).reshape(right_shape) + expected = np.matmul(left, right) * pq.m * pq.s + self.assertQuantityEqual((left * pq.m) @ (right * pq.s), expected) + self.assertQuantityEqual(np.matmul(left * pq.m, right * pq.s), expected) + + def test_plain_array_operands(self): + left = np.array([[1, 2], [3, 4]]) + right = np.array([[5, 6], [7, 8]]) + expected = [[19, 22], [43, 50]] * pq.m + for matmul in (operator.matmul, np.matmul): + for plain_left in (True, False): + with self.subTest(matmul=matmul, plain_left=plain_left): + if plain_left: + result = matmul(left, right * pq.m) + else: + result = matmul(left * pq.m, right) + self.assertQuantityEqual(result, expected) + + def test_complex_vectors(self): + left = Quantity([1 + 2j, 3 - 4j], pq.m) + right = Quantity([5 - 6j, 7 + 8j], pq.s) + # matmul, unlike vdot, does not conjugate the first operand. + self.assertQuantityEqual(left @ right, (70 + 0j) * pq.m * pq.s) + + def test_output_units(self): + left = [[1., 2.], [3., 4.]] * pq.m + right = [[5., 6.], [7., 8.]] * pq.s + out = np.zeros((2, 2)) * pq.kg + self.assertIs(np.matmul(left, right, out=out), out) + self.assertQuantityEqual(out, [[19., 22.], [43., 50.]] * pq.m * pq.s) + self.assertIs(np.matmul(left.magnitude, right.magnitude, out=out), out) + self.assertQuantityEqual(out, [[19., 22.], [43., 50.]] * pq.dimensionless) + + def test_output_view_units(self): + left = [[1., 2.], [3., 4.]] * pq.m + right = [[5., 6.], [7., 8.]] * pq.s + destination = np.zeros((3, 2)) * pq.m * pq.s + out = destination[:2] + self.assertIs(np.matmul(left, right, out=out), out) + self.assertQuantityEqual( + destination, [[19., 22.], [43., 50.], [0., 0.]] * pq.m * pq.s) + + invalid = np.zeros((3, 2)) * pq.kg + with self.assertRaisesRegex(ValueError, 'modify units of a view'): + np.matmul(left, right, out=invalid[:2]) + self.assertEqual(invalid.units, pq.kg) + + @unittest.skipIf(np.lib.NumpyVersion(np.__version__) < '1.25.0', + 'NumPy added in-place matrix multiplication in 1.25') + def test_inplace(self): + left = [[1., 2.], [3., 4.]] * pq.m + right = [[5., 6.], [7., 8.]] * pq.s + left @= right + self.assertQuantityEqual(left, [[19., 22.], [43., 50.]] * pq.m * pq.s) + + def test_invalid_shapes(self): + with self.assertRaises(ValueError): + ([1, 2] * pq.m) @ ([1, 2, 3] * pq.s) + with self.assertRaises(ValueError): + (1 * pq.m) @ ([1, 2] * pq.s) + + def test_uncertainty_remains_unsupported(self): + left = UncertainQuantity([1., 2.], pq.m, uncertainty=[.1, .2]) + right = [3., 4.] * pq.s + for matmul in (operator.matmul, np.matmul): + for args in ((left, right), (right, left)): + with self.subTest(matmul=matmul, args=args): + with self.assertRaises(ValueError): + matmul(*args) + + for out in (Quantity(0., pq.kg), UncertainQuantity(0., pq.kg, uncertainty=1.)): + for args in ((left, right), (right, left)): + with self.subTest(args=args, out_type=type(out)): + with self.assertRaises(ValueError): + np.matmul(*args, out=out) + + out = UncertainQuantity(0., pq.kg, uncertainty=1.) + with self.assertRaises(ValueError): + np.matmul(right, right, out=out) + with self.assertRaises(ValueError): + np.matmul(right.magnitude, right.magnitude, out=out)