Repository: scipy/scipy

Issue:
`binned_statistic_2d` does not accept `values.dtype==complex`
scipy.stats.binned_statistic_2d does not handle complex values, even with custom callables.

The cause seems to be in https://github.com/scipy/scipy/blob/master/scipy/stats/_binned_statistic.py#L546, where the result accumulator gets initialized with a hardcoded float: `result = np.empty([Vdim, nbin.prod()], float)`

### Reproducing code example:
```
import numpy
import scipy.stats

xx = [0.6, 0.7, 1.1]
yy = [.5, .5, .5]
vv = [1j, 1+.3j, 1]
bin_edges = [0, 1, 2]

my_mean = lambda list: numpy.average(list)

binned, _, _, _ = scipy.stats.binned_statistic_2d(xx, yy, vv, statistic='mean', bins=bin_edges)

print(binned)
```

### Error message:
```
Traceback (most recent call last):
  File "scipy_binnedstats_demo.py", line 21, in <module>
    binned, _, _, _ = scipy.stats.binned_statistic_2d(xx, yy, vv, statistic='mean', bins=bin_edges)
  File "\lib\site-packages\scipy\stats\_binned_statistic.py", line 344, in binned_statistic_2d
    expand_binnumbers=expand_binnumbers)
  File "\lib\site-packages\scipy\stats\_binned_statistic.py", line 553, in binned_statistic_dd
    flatsum = np.bincount(binnumbers, values[vv])
TypeError: Cannot cast array data from dtype('complex128') to dtype('float64') according to the rule 'safe'
```
when passing `statistic=my_mean` instead, execution finished, but the result is wrong:
`[[ 0.5  nan] [ 1.   nan]]` - and there is a corresponding `ComplexWarning` as the complex numbers gets cast to real.

### Scipy/Numpy/Python version information:
```
1.0.0 1.13.3 sys.version_info(major=3, minor=6, micro=3, releaselevel='final', serial=0)
```
...but the initialization is still present in ad0b7a1  on Oct 12, 2018

### Partial Fix:
Changing 
```result = np.empty([Vdim, nbin.prod()], float)```
to
```result = np.empty([Vdim, nbin.prod()], values.dtype)``` 
in `binned_statistic_dd()` would fix part of the problem - now, using the custom callable gives the correct results.

I have not submitted this change as a pull request of my own because I came across a problem that requires more thought by someone more familiar with scipy than myself:

It breaks usage with `values.dtype==int` and `statistic='mean'` - which is probably a more common use-case than complex values. Here, an accumulator initialized with `float` is the right choice. 



Discussion:
How about this?

    if np.iscomplexobj(values):
        result = np.empty([Vdim, nbin.prod()], values.dtype)
    else:
        # Use a float64 accumulator to avoid integer overflow
        result = np.empty([Vdim, nbin.prod()], float)

This feature is a bit tricky though. Using a 2-D histogram for a 1-D array of complex values makes sense I think, but shapes are going to be different, and there may be other issues. An n-D histogram for an n-D array doesn't make much sense probably. See e.g. https://stats.stackexchange.com/questions/33582/calculating-and-comparing-histograms-of-complex-numbers. So while we could support this, I don't think we should advertise it.
`np.iscomplexobj` appears to be the right tool - it wasn't part of my numpy vocabulary so far, thanks for bringing it to my attention.

Re: feature trickiness - the function is named `binned_statistic_2d`. It is not "histogram", although that is probably a common use case. More importantly, `statistic='count'` is defined just fine. My question is about: `(x: real, y: real, values: complex)` - you can count/bin any type of data on a 2D (real) grid. If I understand your comment (and the link) correctly, the issue is with complex-valued bin edges. My issue is with real bins containing complex data.

Going through the values for `statistic=` as per documentation:
* The default is `statistic=mean`, which, like `sum`, is unambiguous for complex values. 
* `count` should work, see above.
* `min` and `max` give problems (the clumsy `argmin(abs(z))` is probably the most useful). 
* `median` makes little sense. I have some use-cases for `argmedian(angle(z))`, but they may be rather domain-specific.
* user-defined 'function's require users to think about their data anyways.

To me, enabling complex use with `mean`, `sum` and `function` feels like fixing something. (Re-reading the documentation, I cannot find and hint that the function should not accept complex values Putting a warning in place for use with the sorting functions is probably a good idea, though.
Good points, thanks @ToBeckmann. I agree it would be good to improve support for complex. If you could send a PR, that would be great.



Resolving patch:
diff --git a/scipy/stats/_binned_statistic.py b/scipy/stats/_binned_statistic.py
--- a/scipy/stats/_binned_statistic.py
+++ b/scipy/stats/_binned_statistic.py
@@ -362,6 +362,17 @@ def binned_statistic_2d(x, y, values, statistic='mean',
                                       'binnumber'))
 
 
+def _bincount(x, weights):
+    if np.iscomplexobj(weights):
+        a = np.bincount(x, np.real(weights))
+        b = np.bincount(x, np.imag(weights))
+        z = a + b*1j
+
+    else:
+        z = np.bincount(x, weights)
+    return z
+
+
 def binned_statistic_dd(sample, values, statistic='mean',
                         bins=10, range=None, expand_binnumbers=False,
                         binned_statistic_result=None):
@@ -578,33 +589,39 @@ def binned_statistic_dd(sample, values, statistic='mean',
         dedges = [np.diff(edges[i]) for i in builtins.range(Ndim)]
         binnumbers = binned_statistic_result.binnumber
 
-    result = np.empty([Vdim, nbin.prod()], float)
+    # Avoid overflow with double precision. Complex `values` -> `complex128`.
+    result_type = np.result_type(values, np.float64)
+    result = np.empty([Vdim, nbin.prod()], dtype=result_type)
 
     if statistic in {'mean', np.mean}:
         result.fill(np.nan)
-        flatcount = np.bincount(binnumbers, None)
+        flatcount = _bincount(binnumbers, None)
         a = flatcount.nonzero()
         for vv in builtins.range(Vdim):
-            flatsum = np.bincount(binnumbers, values[vv])
+            flatsum = _bincount(binnumbers, values[vv])
             result[vv, a] = flatsum[a] / flatcount[a]
     elif statistic in {'std', np.std}:
         result.fill(np.nan)
-        flatcount = np.bincount(binnumbers, None)
+        flatcount = _bincount(binnumbers, None)
         a = flatcount.nonzero()
         for vv in builtins.range(Vdim):
-            flatsum = np.bincount(binnumbers, values[vv])
+            flatsum = _bincount(binnumbers, values[vv])
             delta = values[vv] - flatsum[binnumbers] / flatcount[binnumbers]
-            std = np.sqrt(np.bincount(binnumbers, delta**2)[a] / flatcount[a])
+            std = np.sqrt(
+                _bincount(binnumbers, delta*np.conj(delta))[a] / flatcount[a]
+            )
             result[vv, a] = std
+        result = np.real(result)
     elif statistic == 'count':
+        result = np.empty([Vdim, nbin.prod()], dtype=np.float64)
         result.fill(0)
-        flatcount = np.bincount(binnumbers, None)
+        flatcount = _bincount(binnumbers, None)
         a = np.arange(len(flatcount))
         result[:, a] = flatcount[np.newaxis, :]
     elif statistic in {'sum', np.sum}:
         result.fill(0)
         for vv in builtins.range(Vdim):
-            flatsum = np.bincount(binnumbers, values[vv])
+            flatsum = _bincount(binnumbers, values[vv])
             a = np.arange(len(flatsum))
             result[vv, a] = flatsum
     elif statistic in {'median', np.median}:
@@ -635,8 +652,18 @@ def binned_statistic_dd(sample, values, statistic='mean',
                 null = statistic([])
             except Exception:
                 null = np.nan
+        if np.iscomplexobj(null):
+            result = result.astype(np.complex128)
         result.fill(null)
-        _calc_binned_statistic(Vdim, binnumbers, result, values, statistic)
+        try:
+            _calc_binned_statistic(
+                Vdim, binnumbers, result, values, statistic
+            )
+        except ValueError:
+            result = result.astype(np.complex128)
+            _calc_binned_statistic(
+                Vdim, binnumbers, result, values, statistic
+            )
 
     # Shape into a proper matrix
     result = result.reshape(np.append(Vdim, nbin))
@@ -664,7 +691,10 @@ def _calc_binned_statistic(Vdim, bin_numbers, result, values, stat_func):
         bin_map = _create_binned_data(bin_numbers, unique_bin_numbers,
                                       values, vv)
         for i in unique_bin_numbers:
-            result[vv, i] = stat_func(np.array(bin_map[i]))
+            stat = stat_func(np.array(bin_map[i]))
+            if np.iscomplexobj(stat) and not np.iscomplexobj(result):
+                raise ValueError("The statistic function returns complex ")
+            result[vv, i] = stat
 
 
 def _create_binned_data(bin_numbers, unique_bin_numbers, values, vv):


Test patch:
diff --git a/scipy/stats/tests/test_binned_statistic.py b/scipy/stats/tests/test_binned_statistic.py
--- a/scipy/stats/tests/test_binned_statistic.py
+++ b/scipy/stats/tests/test_binned_statistic.py
@@ -1,5 +1,6 @@
 import numpy as np
 from numpy.testing import assert_allclose
+import pytest
 from pytest import raises as assert_raises
 from scipy.stats import (binned_statistic, binned_statistic_2d,
                          binned_statistic_dd)
@@ -535,7 +536,7 @@ def test_gh14332(self):
         size = 20
         for i in range(size):
             x += [1-0.1**i]
- 
+
         bins = np.linspace(0,1,11)
         sum1, edges1, bc = binned_statistic_dd(x, np.ones(len(x)),
                                                bins=[bins], statistic='sum')
@@ -543,3 +544,25 @@ def test_gh14332(self):
 
         assert_allclose(sum1, sum2)
         assert_allclose(edges1[0], edges2)
+
+    @pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+    @pytest.mark.parametrize("statistic", [np.mean, np.median, np.sum, np.std,
+                                           np.min, np.max, 'count',
+                                           lambda x: (x**2).sum(),
+                                           lambda x: (x**2).sum() * 1j])
+    def test_dd_all(self, dtype, statistic):
+        def ref_statistic(x):
+            return len(x) if statistic == 'count' else statistic(x)
+
+        rng = np.random.default_rng(3704743126639371)
+        n = 10
+        x = rng.random(size=n)
+        i = x >= 0.5
+        v = rng.random(size=n)
+        if dtype is np.complex128:
+            v = v + rng.random(size=n)*1j
+
+        stat, _, _ = binned_statistic_dd(x, v, statistic, bins=2)
+        ref = np.array([ref_statistic(v[~i]), ref_statistic(v[i])])
+        assert_allclose(stat, ref)
+        assert stat.dtype == np.result_type(ref.dtype, np.float64)