multiply conj of (x-mean)
@@ -696,6 +696,30 @@ cpdef _ndarray_base _nanstd(_ndarray_base a, axis, dtype, out, ddof, keepdims):
cpdef _ndarray_base _nanvar(_ndarray_base a, axis, dtype, out, ddof, keepdims):
+ dtype_mean = a.dtype
+ dtype_out = numpy.dtype(dtype)
+ if dtype is None:
+ if a.dtype.kind in 'biu':
+ dtype_mean = 'float64'
+ dtype_out = 'float64'
+ else:
+ dtype_mean = a.dtype
+ dtype_out = a.dtype
+ if a.dtype.kind == 'c':
+ dtype_out = numpy.dtype(a.dtype.char.lower())
+
+ if a.dtype.kind == 'c':
+ if out is not None:
+ raise NotImplementedError(
+ 'Variance for complex numbers is not implemented when out != '
+ 'None. Current implemention does not convert the dtype.'
+ )
+ dtype_out = dtype
+ if dtype in ['D', 'F']:
+ dtype == dtype.lower()
+
+ dtype_out = dtype
+
_count = _count_non_nan(a, axis=axis, keepdims=True)
arrsum = _math._nansum(a, axis=axis, dtype=dtype, out=None, keepdims=True)
@@ -703,14 +727,34 @@ cpdef _ndarray_base _nanvar(_ndarray_base a, axis, dtype, out, ddof, keepdims):
return _nanvar_core(
a, arrsum, _count, ddof, axis=axis, keepdims=keepdims)
else:
- return _nanvar_core_out(
+ out = _nanvar_core_out(
a, arrsum, _count, ddof, out, axis=axis, keepdims=keepdims)
+ return out.astype(dtype_out, copy=False)
+
+
+cdef _nanvar_complex = '''
+
+template <typename T>
+__device__ T nanvar_complex_cast(T x) {
+ return x * x;
+}
+
+__device__ float nanvar_complex_cast(const complex<float>& x) {
+ return norm(x);
+}
+
+__device__ double nanvar_complex_cast(const complex<double>& x) {
+ return norm(x);
+}
+
+'''
+
-cdef _nanvar_preamble = '''
+cdef _nanvar_preamble = _nanvar_complex + '''
template <typename S, typename T>
__device__ T nanvar_impl(S x, T mean, long long alpha) {
- return (isnan(x) ? T(0) : T((x - mean) * (x - mean))) / alpha;
+ return (isnan(x) ? T(0) : T(nanvar_complex_cast(x - mean))) / alpha;
}
'''
