complex features (#832)

This PR adds support for

* calling `numpy.{real,imag,conjugate}` with a `Quantity` and
* linearizing a function with complex-valued arguments.

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))

nutils/__init__.py

@@ -1,4 +1,4 @@
 'Numerical Utilities for Finite Element Analysis'
 
-__version__ = version = '9a3'
+__version__ = version = '9a4'
 version_name = 'jook-sing'

nutils/function.py

@@ -1728,7 +1728,7 @@ def linearize(__array: IntoArray, __arguments: Union[str, Dict[str, str], Iterab
     for kv in args:
         k, v = kv.split(':', 1) if isinstance(kv, str) else kv
         f = derivative(array, k)
-        parts.append(numpy.sum(f * Argument(v, f.shape[array.ndim:]), tuple(range(array.ndim, f.ndim))))
+        parts.append(numpy.sum(f * Argument(v, *array.arguments[k]), tuple(range(array.ndim, f.ndim))))
     return util.sum(parts)
 
 

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'))
 

tests/test_function.py

@@ -1377,6 +1377,15 @@ def test(self):
         _q = 5.
         self.assertAllEqual(f.eval(u=_u, v=_v, q=_q).export('dense'), 3 * _u**2 * _v + _q)
 
+    def test_complex(self):
+        f = function.linearize(function.Argument('u', shape=(3, 4), dtype=complex)**3
+                             + function.Argument('p', shape=(), dtype=complex), 'u:v,p:q')
+        # test linearization of u**3 + p -> 3 u**2 v + q through evaluation
+        _u = numpy.array([1+2j, 3+4j, 5+6j])[:,numpy.newaxis].repeat(4, 1)
+        _v = numpy.array([5+1j, 6+2j, 7+3j, 8+4j])[numpy.newaxis,:].repeat(3, 0)
+        _q = 5.
+        self.assertAllEqual(f.eval(u=_u, v=_v, q=_q).export('dense'), 3 * _u**2 * _v + _q)
+
 
 class attributes(TestCase):