diff --git a/dpnp/dpnp_algo/dpnp_arraycreation.py b/dpnp/dpnp_algo/dpnp_arraycreation.py index b6d3c6120683..0399deea254a 100644 --- a/dpnp/dpnp_algo/dpnp_arraycreation.py +++ b/dpnp/dpnp_algo/dpnp_arraycreation.py @@ -1,5 +1,7 @@ +import math import operator +import dpctl.utils as dpu import numpy import dpnp @@ -10,6 +12,7 @@ "dpnp_geomspace", "dpnp_linspace", "dpnp_logspace", + "dpnp_nd_grid", ] @@ -256,3 +259,134 @@ def dpnp_logspace( if dtype is None: return dpnp.power(base, res) return dpnp.power(base, res).astype(dtype, copy=False) + + +class dpnp_nd_grid: + """ + Construct a multi-dimensional "meshgrid". + + ``grid = dpnp_nd_grid()`` creates an instance which will return a mesh-grid + when indexed. The dimension and number of the output arrays are equal + to the number of indexing dimensions. If the step length is not a + complex number, then the stop is not inclusive. + + However, if the step length is a complex number (e.g. 5j), then the + integer part of its magnitude is interpreted as specifying the + number of points to create between the start and stop values, where + the stop value is inclusive. + + If instantiated with an argument of ``sparse=True``, the mesh-grid is + open (or not fleshed out) so that only one-dimension of each returned + argument is greater than 1. + + Parameters + ---------- + sparse : bool, optional + Whether the grid is sparse or not. Default is False. + + """ + + def __init__( + self, sparse=False, device=None, usm_type="device", sycl_queue=None + ): + dpu.validate_usm_type(usm_type, allow_none=False) + self.sparse = sparse + self.usm_type = usm_type + self.sycl_queue_normalized = dpnp.get_normalized_queue_device( + sycl_queue=sycl_queue, device=device + ) + + def __getitem__(self, key): + if isinstance(key, slice): + step = key.step + stop = key.stop + start = key.start + if start is None: + start = 0 + if isinstance(step, complex): + step = abs(step) + length = int(step) + if step != 1: + step = (stop - start) / float(step - 1) + stop = stop + step + return ( + dpnp.arange( + 0, + length, + 1, + dtype=dpnp.default_float_type(), + usm_type=self.usm_type, + sycl_queue=self.sycl_queue_normalized, + ) + * step + + start + ) + else: + return dpnp.arange( + start, + stop, + step, + usm_type=self.usm_type, + sycl_queue=self.sycl_queue_normalized, + ) + + size = [] + dtype = int + for k in range(len(key)): + step = key[k].step + start = key[k].start + stop = key[k].stop + if start is None: + start = 0 + if step is None: + step = 1 + if isinstance(step, complex): + size.append(int(abs(step))) + dtype = dpnp.default_float_type() + else: + size.append( + int(math.ceil((key[k].stop - start) / (step * 1.0))) + ) + if ( + isinstance(step, float) + or isinstance(start, float) + or isinstance(stop, float) + ): + dtype = dpnp.default_float_type() + if self.sparse: + nn = [ + dpnp.arange( + _x, + dtype=_t, + usm_type=self.usm_type, + sycl_queue=self.sycl_queue_normalized, + ) + for _x, _t in zip(size, (dtype,) * len(size)) + ] + else: + nn = dpnp.indices( + size, + dtype, + usm_type=self.usm_type, + sycl_queue=self.sycl_queue_normalized, + ) + for k in range(len(size)): + step = key[k].step + start = key[k].start + stop = key[k].stop + if start is None: + start = 0 + if step is None: + step = 1 + if isinstance(step, complex): + step = int(abs(step)) + if step != 1: + step = (stop - start) / float(step - 1) + nn[k] = nn[k] * step + start + if self.sparse: + slobj = [dpnp.newaxis] * len(size) + for k in range(len(size)): + slobj[k] = slice(None, None) + nn[k] = nn[k][tuple(slobj)] + slobj[k] = dpnp.newaxis + return nn diff --git a/dpnp/dpnp_iface_arraycreation.py b/dpnp/dpnp_iface_arraycreation.py index edcfe7ab3fc7..c006bd2436da 100644 --- a/dpnp/dpnp_iface_arraycreation.py +++ b/dpnp/dpnp_iface_arraycreation.py @@ -52,6 +52,7 @@ dpnp_geomspace, dpnp_linspace, dpnp_logspace, + dpnp_nd_grid, ) __all__ = [ @@ -1452,6 +1453,24 @@ class MGridClass: For full documentation refer to :obj:`numpy.mgrid`. + Parameters + ---------- + device : {None, string, SyclDevice, SyclQueue}, optional + An array API concept of device where the output array is created. + The `device` can be ``None`` (the default), an OneAPI filter selector string, + an instance of :class:`dpctl.SyclDevice` corresponding to a non-partitioned SYCL device, + an instance of :class:`dpctl.SyclQueue`, or a `Device` object returned by + :obj:`dpnp.dpnp_array.dpnp_array.device` property. + usm_type : {"device", "shared", "host"}, optional + The type of SYCL USM allocation for the output array. + sycl_queue : {None, SyclQueue}, optional + A SYCL queue to use for output array allocation and copying. + + Returns + ------- + out : one dpnp.ndarray or tuple of dpnp.ndarray + Returns one array of grid indices, grid.shape = (len(dimensions),) + tuple(dimensions). + Examples -------- >>> import dpnp as np @@ -1466,13 +1485,31 @@ class MGridClass: [0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]]) - >>> np.mgrid[-1:1:5j] + + >>> x = np.mgrid[-1:1:5j] + >>> x array([-1. , -0.5, 0. , 0.5, 1. ]) + >>> x.usm_type + 'device' + + >>> y = np.mgrid(usm_type="host")[-1:1:5j] + >>> y + array([-1. , -0.5, 0. , 0.5, 1. ]) + >>> x.usm_type + 'host' """ def __getitem__(self, key): - return dpnp.array(numpy.mgrid[key]) + return dpnp_nd_grid(sparse=False)[key] + + def __call__(self, device=None, usm_type="device", sycl_queue=None): + return dpnp_nd_grid( + sparse=False, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) mgrid = MGridClass() @@ -1484,23 +1521,56 @@ class OGridClass: For full documentation refer to :obj:`numpy.ogrid`. + Parameters + ---------- + device : {None, string, SyclDevice, SyclQueue}, optional + An array API concept of device where the output array is created. + The `device` can be ``None`` (the default), an OneAPI filter selector string, + an instance of :class:`dpctl.SyclDevice` corresponding to a non-partitioned SYCL device, + an instance of :class:`dpctl.SyclQueue`, or a `Device` object returned by + :obj:`dpnp.dpnp_array.dpnp_array.device` property. + usm_type : {"device", "shared", "host"}, optional + The type of SYCL USM allocation for the output array. + sycl_queue : {None, SyclQueue}, optional + A SYCL queue to use for output array allocation and copying. + + Returns + ------- + out : one dpnp.ndarray or tuple of dpnp.ndarray + Returns a tuple of arrays, with grid[i].shape = (1, ..., 1, dimensions[i], 1, ..., 1) + with dimensions[i] in the ith place. + Examples -------- >>> import dpnp as np - >>> from numpy import ogrid - >>> ogrid[-1:1:5j] - array([-1. , -0.5, 0. , 0.5, 1. ]) - >>> ogrid[0:5,0:5] + >>> np.ogrid[0:5, 0:5] [array([[0], [1], [2], [3], [4]]), array([[0, 1, 2, 3, 4]])] + >>> x = np.ogrid[-1:1:5j] + >>> x + array([-1. , -0.5, 0. , 0.5, 1. ]) + >>> x.usm_type + 'device' + + >>> y = np.ogrid(usm_type="host")[-1:1:5j] + >>> y + array([-1. , -0.5, 0. , 0.5, 1. ]) + >>> x.usm_type + 'host' + """ def __getitem__(self, key): - return dpnp.array(numpy.ogrid[key]) + return dpnp_nd_grid(sparse=True)[key] + + def __call__(self, device=None, usm_type="device", sycl_queue=None): + return dpnp_nd_grid( + sparse=True, device=device, usm_type=usm_type, sycl_queue=sycl_queue + ) ogrid = OGridClass() diff --git a/tests/skipped_tests.tbl b/tests/skipped_tests.tbl index 87a29e2cb0c5..31de31445645 100644 --- a/tests/skipped_tests.tbl +++ b/tests/skipped_tests.tbl @@ -177,10 +177,6 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy 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/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid3 -tests/third_party/cupy/creation_tests/test_ranges.py::TestMgrid::test_mgrid3 -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/indexing_tests/test_generate.py::TestAxisConcatenator::test_AxisConcatenator_init1 tests/third_party/cupy/indexing_tests/test_generate.py::TestAxisConcatenator::test_len tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_1 diff --git a/tests/skipped_tests_gpu.tbl b/tests/skipped_tests_gpu.tbl index b6f6ceb45913..19a129c036a8 100644 --- a/tests/skipped_tests_gpu.tbl +++ b/tests/skipped_tests_gpu.tbl @@ -255,11 +255,6 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy 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/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid3 -tests/third_party/cupy/creation_tests/test_ranges.py::TestMgrid::test_mgrid3 -tests/third_party/cupy/creation_tests/test_ranges.py::TestMgrid::test_mgrid5 -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_arange_negative_size tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_no_dtype_int diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index 815a131b5e8d..c2173c887c46 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -1359,6 +1359,18 @@ def test_indices(device, sparse): assert i.sycl_device == device +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +@pytest.mark.parametrize("func", ["mgrid", "ogrid"]) +def test_grid(device, func): + sycl_queue = dpctl.SyclQueue(device) + x = getattr(dpnp, func)(sycl_queue=sycl_queue)[0:4] + assert_sycl_queue_equal(x.sycl_queue, sycl_queue) + + @pytest.mark.parametrize( "device", valid_devices, diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 42cde3e7412d..d8f5640040f8 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -488,6 +488,18 @@ def test_take(func, usm_type_x, usm_type_ind): assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_ind]) +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_indices(usm_type): + x = dp.indices((2,), usm_type=usm_type) + assert x.usm_type == usm_type + + +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +@pytest.mark.parametrize("func", ["mgrid", "ogrid"]) +def test_grid(usm_type, func): + assert getattr(dp, func)(usm_type=usm_type)[0:4].usm_type == usm_type + + @pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) @pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"]) def test_indices(usm_type, sparse): diff --git a/tests/third_party/cupy/creation_tests/test_ranges.py b/tests/third_party/cupy/creation_tests/test_ranges.py index 623adc409b7e..83170cb3e37f 100644 --- a/tests/third_party/cupy/creation_tests/test_ranges.py +++ b/tests/third_party/cupy/creation_tests/test_ranges.py @@ -363,7 +363,7 @@ def test_mgrid1(self, xp): def test_mgrid2(self, xp): return xp.mgrid[-10:10:10j] - @testing.numpy_cupy_array_equal() + @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64()) def test_mgrid3(self, xp): x = xp.zeros(10)[:, None] y = xp.ones(10)[:, None] @@ -374,7 +374,7 @@ def test_mgrid4(self, xp): # check len(keys) > 1 return xp.mgrid[-10:10:10j, -10:10:10j] - @testing.numpy_cupy_array_equal() + @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64()) def test_mgrid5(self, xp): # check len(keys) > 1 x = xp.zeros(10)[:, None] @@ -396,18 +396,18 @@ def test_ogrid1(self, xp): def test_ogrid2(self, xp): return xp.ogrid[-10:10:10j] - @testing.numpy_cupy_array_equal() + @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64()) def test_ogrid3(self, xp): x = xp.zeros(10)[:, None] y = xp.ones(10)[:, None] return xp.ogrid[x:y:10j] - @testing.numpy_cupy_array_equal() + @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64()) def test_ogrid4(self, xp): # check len(keys) > 1 return xp.ogrid[-10:10:10j, -10:10:10j] - @testing.numpy_cupy_array_equal() + @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64()) def test_ogrid5(self, xp): # check len(keys) > 1 x = xp.zeros(10)[:, None]