add dpnp.nanvar() (#720)

Co-authored-by: Sergey Shalnov <[email protected]>

Author: Rubtsowa

dpnp/backend/include/dpnp_iface.hpp

@@ -189,6 +189,18 @@ template <typename _DataType>
 INP_DLLEXPORT void
     dpnp_matmul_c(void* array1, void* array2, void* result1, size_t size_m, size_t size_n, size_t size_k);
 
+/**
+ * @ingroup BACKEND_API
+ * @brief Compute the variance along the specified axis, while ignoring NaNs.
+ *
+ * @param [in]  array     Input array.
+ * @param [in]  mask_arr  Input mask array when elem is nan.
+ * @param [out] result    Output array.
+ * @param [in]  size      Number of elements in input arrays.
+ */
+template <typename _DataType>
+INP_DLLEXPORT void dpnp_nanvar_c(void* array, void* mask_arr, void* result, size_t size);
+
 /**
  * @ingroup BACKEND_API
  * @brief Return the indices of the elements that are non-zero.

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -135,6 +135,7 @@ enum class DPNPFuncName : size_t
     DPNP_FN_MINIMUM,                  /**< Used in numpy.minimum() implementation  */
     DPNP_FN_MODF,                     /**< Used in numpy.modf() implementation  */
     DPNP_FN_MULTIPLY,                 /**< Used in numpy.multiply() implementation  */
+    DPNP_FN_NANVAR,                   /**< Used in numpy.nanvar() implementation  */
     DPNP_FN_NEGATIVE,                 /**< Used in numpy.negative() implementation  */
     DPNP_FN_NONZERO,                  /**< Used in numpy.nonzero() implementation  */
     DPNP_FN_ONES,                     /**< Used in numpy.ones() implementation */

dpnp/backend/kernels/dpnp_krnl_statistics.cpp

@@ -447,6 +447,36 @@ void dpnp_min_c(void* array1_in, void* result1, const size_t* shape, size_t ndim
     return;
 }
 
+template <typename _DataType>
+void dpnp_nanvar_c(void* array1_in, void* mask_arr1, void* result1, size_t arr_size)
+{
+    _DataType* array1 = reinterpret_cast<_DataType*>(array1_in);
+    bool* mask_arr = reinterpret_cast<bool*>(mask_arr1);
+    _DataType* result = reinterpret_cast<_DataType*>(result1);
+
+    if ((array1 == nullptr) || (mask_arr == nullptr) || (result == nullptr))
+    {
+        return;
+    }
+
+    if (arr_size == 0)
+    {
+        return;
+    }
+
+    size_t ind = 0;
+    for (size_t i = 0; i < arr_size; ++i)
+    {
+        if (!mask_arr[i])
+        {
+            result[ind] = array1[i];
+            ind += 1;
+        }
+    }
+
+    return;
+}
+
 template <typename _DataType, typename _ResultType>
 void dpnp_std_c(
     void* array1_in, void* result1, const size_t* shape, size_t ndim, const size_t* axis, size_t naxis, size_t ddof)
@@ -560,6 +590,11 @@ void func_map_init_statistics(func_map_t& fmap)
     fmap[DPNPFuncName::DPNP_FN_MIN][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_min_c<float>};
     fmap[DPNPFuncName::DPNP_FN_MIN][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_min_c<double>};
 
+    fmap[DPNPFuncName::DPNP_FN_NANVAR][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_nanvar_c<int>};
+    fmap[DPNPFuncName::DPNP_FN_NANVAR][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_nanvar_c<long>};
+    fmap[DPNPFuncName::DPNP_FN_NANVAR][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_nanvar_c<float>};
+    fmap[DPNPFuncName::DPNP_FN_NANVAR][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_nanvar_c<double>};
+
     fmap[DPNPFuncName::DPNP_FN_STD][eft_INT][eft_INT] = {eft_DBL, (void*)dpnp_std_c<int, double>};
     fmap[DPNPFuncName::DPNP_FN_STD][eft_LNG][eft_LNG] = {eft_DBL, (void*)dpnp_std_c<long, double>};
     fmap[DPNPFuncName::DPNP_FN_STD][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_std_c<float, float>};

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -108,6 +108,7 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_MINIMUM
         DPNP_FN_MODF
         DPNP_FN_MULTIPLY
+        DPNP_FN_NANVAR
         DPNP_FN_NEGATIVE
         DPNP_FN_NONZERO
         DPNP_FN_ONES

dpnp/dpnp_algo/dpnp_algo_statistics.pyx

@@ -48,13 +48,15 @@ __all__ += [
     "dpnp_mean",
     "dpnp_median",
     "dpnp_min",
+    "dpnp_nanvar",
     "dpnp_std",
     "dpnp_var",
 ]
 
 
 # C function pointer to the C library template functions
 ctypedef void(*fptr_custom_cov_1in_1out_t)(void * , void * , size_t, size_t)
+ctypedef void(*fptr_custom_nanvar_t)(void * , void * , void * , size_t)
 ctypedef void(*fptr_custom_std_var_1in_1out_t)(void * , void * , size_t * , size_t, size_t * , size_t, size_t)
 
 # C function pointer to the C library template functions
@@ -426,6 +428,25 @@ cpdef dparray dpnp_min(dparray input, axis):
     return _dpnp_min(input, axis_, output_shape)
 
 
+cpdef dparray dpnp_nanvar(dparray arr, ddof):
+    cdef dparray mask_arr = dpnp.isnan(arr)
+    n = sum(mask_arr)
+    res_size = arr.size - n
+
+    cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(arr.dtype)
+
+    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_NANVAR, param1_type, param1_type)
+
+    result_type = dpnp_DPNPFuncType_to_dtype( < size_t > kernel_data.return_type)
+    cdef dparray without_nan_arr = dparray((res_size, ), dtype=result_type)
+
+    cdef fptr_custom_nanvar_t func = <fptr_custom_nanvar_t > kernel_data.ptr
+
+    func(arr.get_data(), mask_arr.get_data(), without_nan_arr.get_data(), arr.size)
+
+    return call_fptr_custom_std_var_1in_1out(DPNP_FN_VAR, without_nan_arr, ddof)
+
+
 cpdef dparray dpnp_std(dparray a, size_t ddof):
     return call_fptr_custom_std_var_1in_1out(DPNP_FN_STD, a, ddof)
 

dpnp/dpnp_iface_statistics.py

@@ -58,6 +58,7 @@
     'mean',
     'median',
     'min',
+    'nanvar',
     'std',
     'var',
 ]
@@ -474,6 +475,44 @@ def min(input, axis=None, out=None, keepdims=numpy._NoValue, initial=numpy._NoVa
     return call_origin(numpy.min, input, axis, out, keepdims, initial, where)
 
 
+def nanvar(arr, axis=None, dtype=None, out=None, ddof=0, keepdims=numpy._NoValue):
+    """
+    Compute the variance along the specified axis, while ignoring NaNs.
+
+    For full documentation refer to :obj:`numpy.nanvar`.
+
+    Limitations
+    -----------
+    Input array is supported as :obj:`dpnp.ndarray`.
+    Prameters ``axis`` is supported only with default value ``None``.
+    Prameters ``dtype`` is supported only with default value ``None``.
+    Prameters ``out`` is supported only with default value ``None``.
+    Prameters ``keepdims`` is supported only with default value ``numpy._NoValue``.
+    Otherwise the function will be executed sequentially on CPU.
+    """
+    if not use_origin_backend(arr):
+        if not isinstance(arr, dparray):
+            pass
+        elif axis is not None:
+            pass
+        elif dtype is not None:
+            pass
+        elif out is not None:
+            pass
+        elif keepdims is not numpy._NoValue:
+            pass
+        else:
+            result = dpnp_nanvar(arr, ddof)
+
+            # scalar returned
+            if result.shape == (1,):
+                return result.dtype.type(result[0])
+
+            return result
+
+    return call_origin(numpy.nanvar, arr, axis=axis, dtype=dtype, out=out, ddof=ddof, keepdims=keepdims)
+
+
 def std(a, axis=None, dtype=None, out=None, ddof=0, keepdims=numpy._NoValue):
     """
     Compute the standard deviation along the specified axis.

tests/test_statistics.py

@@ -18,3 +18,45 @@ def test_median(type, size):
     dpnp_res = dpnp.median(ia)
 
     numpy.testing.assert_allclose(dpnp_res, np_res)
+
+
+@pytest.mark.parametrize("array",
+                         [[2, 0, 6, 2],
+                          [2, 0, 6, 2, 5, 6, 7, 8],
+                          [],
+                          [2, 1, numpy.nan, 5, 3],
+                          [-1, numpy.nan, 1, numpy.inf],
+                          [3, 6, 0, 1],
+                          [3, 6, 0, 1, 8],
+                          [3, 2, 9, 6, numpy.nan],
+                          [numpy.nan, numpy.nan, numpy.inf, numpy.nan],
+                          [[2, 0], [6, 2]],
+                          [[2, 0, 6, 2], [5, 6, 7, 8]],
+                          [[[2, 0], [6, 2]], [[5, 6], [7, 8]]],
+                          [[-1, numpy.nan], [1, numpy.inf]],
+                          [[numpy.nan, numpy.nan], [numpy.inf, numpy.nan]]],
+                         ids=['[2, 0, 6, 2]',
+                              '[2, 0, 6, 2, 5, 6, 7, 8]',
+                              '[]',
+                              '[2, 1, np.nan, 5, 3]',
+                              '[-1, np.nan, 1, np.inf]',
+                              '[3, 6, 0, 1]',
+                              '[3, 6, 0, 1, 8]',
+                              '[3, 2, 9, 6, np.nan]',
+                              '[np.nan, np.nan, np.inf, np.nan]',
+                              '[[2, 0], [6, 2]]',
+                              '[[2, 0, 6, 2], [5, 6, 7, 8]]',
+                              '[[[2, 0], [6, 2]], [[5, 6], [7, 8]]]',
+                              '[[-1, np.nan], [1, np.inf]]',
+                              '[[np.nan, np.nan], [np.inf, np.nan]]'])
+def test_nanvar(array):
+    a = numpy.array(array)
+    ia = dpnp.array(a)
+    for ddof in range(a.ndim):
+        expected = numpy.nanvar(a, ddof=ddof)
+        result = dpnp.nanvar(ia, ddof=ddof)
+        numpy.testing.assert_array_equal(expected, result)
+
+    expected = numpy.nanvar(a, axis=None, ddof=0)
+    result = dpnp.nanvar(ia, axis=None, ddof=0)
+    numpy.testing.assert_array_equal(expected, result)