From e040cd792a7daa389511049a305caeae362c106d Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Tue, 5 Sep 2023 09:57:35 +0200
Subject: [PATCH 1/2] support numpy.{real,imag,conjugate} on Quantities

This patch adds support for calling `numpy.{real,imag,conjugate}` with a
`Quantity`.
---
 nutils/SI.py     | 2 +-
 tests/test_SI.py | 9 +++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/nutils/SI.py b/nutils/SI.py
index 4124cad96..d0e0a627e 100644
--- 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
index 5f159db50..7c42e22e0 100644
--- 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'))
 

From 4c9eb7202f30129738e381c18201189c951c6b40 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Tue, 5 Sep 2023 10:15:13 +0200
Subject: [PATCH 2/2] support complex arguments in function.linearize

This patch adds support for linearizing functions with complex-valued
arguments by instantiating `function.Argument`s with the correct dtype
in `function.linearize`.
---
 nutils/function.py     | 2 +-
 tests/test_function.py | 9 +++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/nutils/function.py b/nutils/function.py
index 7dc89a5df..97497498b 100644
--- a/nutils/function.py
+++ b/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)
 
 
diff --git a/tests/test_function.py b/tests/test_function.py
index fbc5c61cb..4176d26a0 100644
--- a/tests/test_function.py
+++ b/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):