for complex arrays all NaN values are considered equivalent

numpy · charris · Feb 12, 2021 · Dec 25, 2020 · Dec 25, 2020 · Dec 25, 2020
commit f2a99c604a9e1f4d11320ca03c8f55f6bc6eb85b
diff --git a/doc/release/upcoming_changes/18070.improvement.rst b/doc/release/upcoming_changes/18070.improvement.rst
@@ -3,4 +3,10 @@
 When ``np.unique`` operated on an array with multiple ``NaN`` entries,
 its return included a ``NaN`` for each entry that was ``NaN`` in the original array.
 This is now improved such that the returned array contains just one ``NaN`` as the
-last element.
+last element.
+
+Also for complex arrays all ``NaN`` values are considered equivalent
+(no matter whether the ``NaN`` is in the real or imaginary part). As the
+representant for the returned array the smallest one in the
+lexicographical order is chosen - see ``np.sort`` for how the lexicographical
+order is defined for complex arrays.
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py
@@ -213,6 +213,12 @@ def unique(ar, return_index=False, return_inverse=False,
         If nan values are in the input array, a single nan is put
         to the end of the sorted unique values.
 
+        Also for complex arrays all NaN values are considered equivalent
+        (no matter whether the NaN is in the real or imaginary part).
+        As the representant for the returned array the smallest one in the
+        lexicographical order is chosen - see np.sort for how the lexicographical
+        order is defined for complex arrays.
+
        Examples
     --------
     >>> np.unique([1, 1, 2, 2, 3, 3])
@@ -329,7 +335,10 @@ def _unique1d(ar, return_index=False, return_inverse=False,
     mask = np.empty(aux.shape, dtype=np.bool_)
     mask[:1] = True
     if aux.shape[0] > 0 and aux.dtype.kind in "cfmM" and np.isnan(aux[-1]):
-        aux_firstnan = np.searchsorted(aux, aux[-1], side='left')
+        if aux.dtype.kind == "c":  # for complex all NaNs are considered equivalent
+            aux_firstnan = np.searchsorted(np.isnan(aux), True, side='left')
+        else:
+            aux_firstnan = np.searchsorted(aux, aux[-1], side='left')
         mask[1:aux_firstnan] = (aux[1:aux_firstnan] != aux[:aux_firstnan - 1])
         mask[aux_firstnan] = True
         mask[aux_firstnan + 1:] = False

diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py
@@ -542,6 +542,17 @@ def check_all(a, b, i1, i2, c, dt):
         assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv))
         assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt))
 
+        # test for ticket 2111 - complex
+        a = [2.0-1j, np.nan, 1.0+1j, complex(0.0, np.nan), complex(1.0, np.nan)]
+        ua = [1.0+1j, 2.0-1j, complex(0.0, np.nan)]
+        ua_idx = [2, 0, 3]
+        ua_inv = [1, 2, 0, 2, 2]
+        ua_cnt = [1, 1, 3]
+        assert_equal(np.unique(a), ua)
+        assert_equal(np.unique(a, return_index=True), (ua, ua_idx))
+        assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv))
+        assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt))
+
         # test for ticket 2111 - datetime64
         nat = np.datetime64('nat')
         a = [np.datetime64('2020-12-26'), nat, np.datetime64('2020-12-24'), nat]