Meshgrid func (#337)

* meshgrid impl

Author: Alexander-Makaryev

dpnp/backend_arraycreation.pyx

@@ -45,6 +45,7 @@ __all__ += [
     "dpnp_geomspace",
     "dpnp_linspace",
     "dpnp_logspace",
+    "dpnp_meshgrid",
     "dpnp_tri",
     "dpnp_tril",
     "dpnp_triu",
@@ -141,6 +142,54 @@ cpdef dparray dpnp_logspace(start, stop, num, endpoint, base, dtype, axis):
     return dpnp.power(base, temp).astype(dtype)
 
 
+cpdef list dpnp_meshgrid(xi, copy, sparse, indexing):
+    cdef dparray res_item
+
+    input_count = len(xi)
+
+    # simple case
+    if input_count == 0:
+        return []
+
+    # simple case
+    if input_count == 1:
+        return [dpnp.copy(xi[0])]
+
+    shape_mult = 1
+    for i in range(input_count):
+        shape_mult = shape_mult * xi[i].size
+
+    shape_list = []
+    for i in range(input_count):
+        shape_list.append(xi[i].size)
+    if indexing == "xy":
+        temp = shape_list[0]
+        shape_list[0] = shape_list[1]
+        shape_list[1] = temp
+
+    steps = []
+    for i in range(input_count):
+        shape_mult = shape_mult // shape_list[i]
+        steps.append(shape_mult)
+    if indexing == "xy":
+        temp = steps[0]
+        steps[0] = steps[1]
+        steps[1] = temp
+
+    shape = tuple(shape_list)
+
+    result = []
+    for i in range(input_count):
+        res_item = dparray(shape=shape, dtype=xi[i].dtype)
+
+        for j in range(res_item.size):
+            res_item[j] = xi[i][(j // steps[i]) % xi[i].size]
+
+        result.append(res_item)
+
+    return result
+
+
 cpdef dparray dpnp_tri(N, M, k, dtype):
     cdef dparray result
 

dpnp/dpnp_iface_arraycreation.py

@@ -63,6 +63,7 @@
     "geomspace",
     "linspace",
     "logspace",
+    "meshgrid",
     "ones",
     "ones_like",
     "tri",
@@ -354,6 +355,7 @@ def diag(v, k=0):
 
     Examples
     --------
+    >>> import dpnp as np
     >>> x = np.arange(9).reshape((3,3))
     >>> x
     array([[0, 1, 2],
@@ -391,6 +393,7 @@ def diagflat(v, k=0):
 
     Examples
     --------
+    >>> import dpnp as np
     >>> np.diagflat([[1,2], [3,4]])
     array([[1, 0, 0, 0],
            [0, 2, 0, 0],
@@ -706,6 +709,68 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
     return call_origin(numpy.logspace, start, stop, num, endpoint, base, dtype, axis)
 
 
+def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
+    """
+    Return coordinate matrices from coordinate vectors.
+
+    Make N-D coordinate arrays for vectorized evaluations of
+    N-D scalar/vector fields over N-D grids, given
+    one-dimensional coordinate arrays x1, x2,..., xn.
+
+    For full documentation refer to :obj:`numpy.meshgrid`.
+
+    Limitations
+    -----------
+    Parameter ``copy`` is supported only with default value ``True``.
+    Parameter ``sparse`` is supported only with default value ``False``.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> nx, ny = (3, 2)
+    >>> x = np.linspace(0, 1, nx)
+    >>> y = np.linspace(0, 1, ny)
+    >>> xv, yv = np.meshgrid(x, y)
+    >>> xv
+    array([[0. , 0.5, 1. ],
+           [0. , 0.5, 1. ]])
+    >>> yv
+    array([[0.,  0.,  0.],
+           [1.,  1.,  1.]])
+    >>> xv, yv = np.meshgrid(x, y, sparse=True)  # make sparse output arrays
+    >>> xv
+    array([[0. ,  0.5,  1. ]])
+    >>> yv
+    array([[0.],
+           [1.]])
+
+    `meshgrid` is very useful to evaluate functions on a grid.
+
+    >>> import matplotlib.pyplot as plt
+    >>> x = np.arange(-5, 5, 0.1)
+    >>> y = np.arange(-5, 5, 0.1)
+    >>> xx, yy = np.meshgrid(x, y, sparse=True)
+    >>> z = np.sin(xx**2 + yy**2) / (xx**2 + yy**2)
+    >>> h = plt.contourf(x,y,z)
+    >>> plt.show()
+
+    """
+
+    if not use_origin_backend():
+        # original limitation
+        if indexing not in ["ij", "xy"]:
+            checker_throw_value_error("meshgrid", "indexing", indexing, "'ij' or 'xy'")
+
+        if copy is not True:
+            checker_throw_value_error("meshgrid", "copy", copy, True)
+        if sparse is not False:
+            checker_throw_value_error("meshgrid", "sparse", sparse, False)
+
+        return dpnp_meshgrid(xi, copy, sparse, indexing)
+
+    return call_origin(numpy.meshgrid, xi, copy, sparse, indexing)
+
+
 def ones(shape, dtype=None, order='C'):
     """
     Return a new array of given shape and type, filled with ones.
@@ -797,6 +862,7 @@ def tri(N, M=None, k=0, dtype=numpy.float):
 
     Examples
     --------
+    >>> import dpnp as np
     >>> np.tri(3, 5, 2, dtype=int)
     array([[1, 1, 1, 0, 0],
            [1, 1, 1, 1, 0],
@@ -825,6 +891,7 @@ def tril(m, k=0):
 
     Examples
     --------
+    >>> import dpnp as np
     >>> np.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)
     array([[ 0,  0,  0],
            [ 4,  0,  0],
@@ -853,6 +920,7 @@ def triu(m, k=0):
 
     Examples
     --------
+    >>> import dpnp as np
     >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1)
     array([[ 1,  2,  3],
            [ 4,  5,  6],

tests/skipped_tests.tbl

@@ -315,14 +315,6 @@ tests/third_party/cupy/creation_tests/test_from_data.py::TestArrayCopy_param_12_
 tests/third_party/cupy/creation_tests/test_from_data.py::TestArrayCopy_param_13_{copy=False, ndmin=2, xp=dpnp}::test_cupy_array
 tests/third_party/cupy/creation_tests/test_from_data.py::TestArrayCopy_param_14_{copy=False, ndmin=3, xp=numpy}::test_cupy_array
 tests/third_party/cupy/creation_tests/test_from_data.py::TestArrayCopy_param_15_{copy=False, ndmin=3, xp=dpnp}::test_cupy_array
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange4
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange5
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange8
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange9
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_negative_size
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_no_dtype_float
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_no_dtype_int
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_linspace_array_start_stop
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_linspace_array_start_stop_axis1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_linspace_float_underflow
@@ -357,18 +349,10 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid2
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid3
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid0
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid2
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid3
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid0
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid2

tests/skipped_tests_gpu.tbl

@@ -365,18 +365,10 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid2
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid3
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid0
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid2
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid3
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid0
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid2
@@ -393,14 +385,6 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid2
 tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid3
 tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid4
 tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid5
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange4
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange5
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange8
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange9
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_negative_size
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_no_dtype_float
-tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_no_dtype_int
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_linspace_array_start_stop
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_linspace_array_start_stop_axis1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_linspace_float_underflow