support numpy.{real,imag,conjugate} on Quantities

This patch adds support for calling `numpy.{real,imag,conjugate}` with a `Quantity`.
evalf · Sep 5, 2023 · e040cd7 · e040cd7
1 parent f28fb33
commit e040cd7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/nutils/SI.py b/nutils/SI.py
@@ -326,7 +326,7 @@ def register(*names, __table=__DISPATCH_TABLE):
               'trace', 'ptp', 'amax', 'amin', 'max', 'min', 'mean', 'take',
               'broadcast_to', 'transpose', 'getitem', 'opposite', 'jump',
               'replace_arguments', 'linearize', 'derivative', 'integral',
-              'sample', 'scatter', 'kronecker')
+              'sample', 'scatter', 'kronecker', 'real', 'imag', 'conjugate')
     def __unary(op, *args, **kwargs):
         (dim0, arg0), = Quantity.__unpack(args[0])
         return dim0.wrap(op(arg0, *args[1:], **kwargs))

diff --git a/tests/test_SI.py b/tests/test_SI.py
@@ -129,6 +129,15 @@ def test_pos(self):
     def test_abs(self):
         self.assertEqual(numpy.abs(SI.Mass('-2kg')), SI.Mass('2kg'))
 
+    def test_real(self):
+        self.assertEqual(numpy.real(SI.ElectricPotential('1V') + 1j * SI.ElectricPotential('2V')), SI.ElectricPotential('1V'))
+
+    def test_imag(self):
+        self.assertEqual(numpy.imag(SI.ElectricPotential('1V') + 1j * SI.ElectricPotential('2V')), SI.ElectricPotential('2V'))
+
+    def test_conjugate(self):
+        self.assertEqual(numpy.conjugate(SI.ElectricPotential('1V') + 1j * SI.ElectricPotential('2V')), SI.ElectricPotential('1V') - 1j * SI.ElectricPotential('2V'))
+
     def test_sqrt(self):
         self.assertEqual(numpy.sqrt(SI.Area('4m2')), SI.Length('2m'))