Repository: numpy/numpy

Issue:
PERF: Variance of complex array
```python
import sys
import numpy as np

print(sys.version)
# 3.7.4 (default, Aug 13 2019, 15:17:50) 
# [Clang 4.0.1 (tags/RELEASE_401/final)]
print(np.__version__)
# 1.18.1

np.random.seed(2357)
n = 10 ** 8
arr = np.random.randn(n) + 1j * np.random.randn(n)

%timeit arr.var()
# 1.99 s ± 36 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit arr.real.var() + arr.imag.var()
# 1.25 s ± 2.12 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

v1 = arr.var()
v2 = arr.real.var() + arr.imag.var()

np.testing.assert_almost_equal(v1, v2)
```

It looks like summing the variances of the real and imaginary parts of a complex array is almost twice as fast as using `var` directly (and gives the same answer), so maybe there's something slowing down the calculation in this case (or maybe it makes sense to calculate the variance in this way)?


Discussion:
I see the same thing on 2fd9ff8 though to a lesser degree:
```python
>>> %timeit arr.var()
863 ms ± 7.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit arr.real.var() + arr.imag.var()
713 ms ± 2.49 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

A closer look with `line_profiler` IDs the line calling `conjugate` as the likely culprit:

https://github.com/numpy/numpy/blob/6894bbc6d396b87464cbc21516d239d5f94f13b7/numpy/core/_methods.py#L191-L195

Resolving patch:
diff --git a/benchmarks/benchmarks/bench_core.py b/benchmarks/benchmarks/bench_core.py
--- a/benchmarks/benchmarks/bench_core.py
+++ b/benchmarks/benchmarks/bench_core.py
@@ -180,3 +180,14 @@ def time_unpackbits_axis1_little(self):
 class Indices(Benchmark):
     def time_indices(self):
         np.indices((1000, 500))
+
+class VarComplex(Benchmark):
+    params = [10**n for n in range(1, 9)]
+    def setup(self, n):
+        self.arr = np.random.randn(n) + 1j * np.random.randn(n)
+
+    def teardown(self, n):
+        del self.arr
+
+    def time_var(self, n):
+        self.arr.var()
diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py
--- a/numpy/core/_methods.py
+++ b/numpy/core/_methods.py
@@ -21,6 +21,21 @@
 umr_any = um.logical_or.reduce
 umr_all = um.logical_and.reduce
 
+# Complex types to -> (2,)float view for fast-path computation in _var()
+_complex_to_float = {
+    nt.dtype(nt.csingle) : nt.dtype(nt.single),
+    nt.dtype(nt.cdouble) : nt.dtype(nt.double),
+}
+# Special case for windows: ensure double takes precedence
+if nt.dtype(nt.longdouble) != nt.dtype(nt.double):
+    _complex_to_float.update({
+        nt.dtype(nt.clongdouble) : nt.dtype(nt.longdouble),
+    })
+# Add reverse-endian types
+_complex_to_float.update({
+    k.newbyteorder() : v.newbyteorder() for k, v in _complex_to_float.items()
+})
+
 # avoid keyword arguments to speed up parsing, saves about 15%-20% for very
 # small reductions
 def _amax(a, axis=None, out=None, keepdims=False,
@@ -189,8 +204,16 @@ def _var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False):
     # Note that x may not be inexact and that we need it to be an array,
     # not a scalar.
     x = asanyarray(arr - arrmean)
+
     if issubclass(arr.dtype.type, (nt.floating, nt.integer)):
         x = um.multiply(x, x, out=x)
+    # Fast-paths for built-in complex types
+    elif x.dtype in _complex_to_float:
+        xv = x.view(dtype=(_complex_to_float[x.dtype], (2,)))
+        um.multiply(xv, xv, out=xv)
+        x = um.add(xv[..., 0], xv[..., 1], out=x.real).real
+    # Most general case; includes handling object arrays containing imaginary
+    # numbers and complex types with non-native byteorder
     else:
         x = um.multiply(x, um.conjugate(x), out=x).real
 


Test patch:
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -5507,6 +5507,39 @@ def test_var_values(self):
                 res = _var(mat, axis=axis)
                 assert_almost_equal(res, tgt)
 
+    @pytest.mark.parametrize(('complex_dtype', 'ndec'), (
+        ('complex64', 6),
+        ('complex128', 7),
+        ('clongdouble', 7),
+    ))
+    def test_var_complex_values(self, complex_dtype, ndec):
+        # Test fast-paths for every builtin complex type
+        for axis in [0, 1, None]:
+            mat = self.cmat.copy().astype(complex_dtype)
+            msqr = _mean(mat * mat.conj(), axis=axis)
+            mean = _mean(mat, axis=axis)
+            tgt = msqr - mean * mean.conjugate()
+            res = _var(mat, axis=axis)
+            assert_almost_equal(res, tgt, decimal=ndec)
+
+    def test_var_dimensions(self):
+        # _var paths for complex number introduce additions on views that
+        # increase dimensions. Ensure this generalizes to higher dims
+        mat = np.stack([self.cmat]*3)
+        for axis in [0, 1, 2, -1, None]:
+            msqr = _mean(mat * mat.conj(), axis=axis)
+            mean = _mean(mat, axis=axis)
+            tgt = msqr - mean * mean.conjugate()
+            res = _var(mat, axis=axis)
+            assert_almost_equal(res, tgt)
+
+    def test_var_complex_byteorder(self):
+        # Test that var fast-path does not cause failures for complex arrays
+        # with non-native byteorder
+        cmat = self.cmat.copy().astype('complex128')
+        cmat_swapped = cmat.astype(cmat.dtype.newbyteorder())
+        assert_almost_equal(cmat.var(), cmat_swapped.var())
+
     def test_std_values(self):
         for mat in [self.rmat, self.cmat, self.omat]:
             for axis in [0, 1, None]: