test_ranges.py

import functools
import math
import sys
import unittest

import numpy
import pytest

import dpnp as cupy
from tests.helper import has_support_aspect64
from tests.third_party.cupy import testing


def skip_int_equality_before_numpy_1_20(names=("dtype",)):
    """Require numpy/numpy#16841 or skip the equality check."""

    def decorator(wrapped):
        if numpy.lib.NumpyVersion(numpy.__version__) >= "1.20.0":
            return wrapped

        @functools.wraps(wrapped)
        def wrapper(self, *args, **kwargs):
            xp = kwargs["xp"]
            dtypes = [kwargs[name] for name in names]
            ret = wrapped(self, *args, **kwargs)
            if any(numpy.issubdtype(dtype, numpy.integer) for dtype in dtypes):
                ret = xp.zeros_like(ret)
            return ret

        return wrapper

    return decorator


class TestRanges(unittest.TestCase):
    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_arange(self, xp, dtype):
        return xp.arange(10, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_arange2(self, xp, dtype):
        return xp.arange(5, 10, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_arange3(self, xp, dtype):
        return xp.arange(1, 11, 2, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_arange4(self, xp, dtype):
        return xp.arange(20, 2, -3, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_arange5(self, xp, dtype):
        return xp.arange(0, 100, None, dtype=dtype)

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_arange6(self, xp, dtype):
        return xp.arange(0, 2, dtype=dtype)

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_arange7(self, xp, dtype):
        return xp.arange(10, 11, dtype=dtype)

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_arange8(self, xp, dtype):
        return xp.arange(10, 8, -1, dtype=dtype)

    def test_arange9(self):
        for xp in (numpy, cupy):
            with pytest.raises((ValueError, TypeError)):
                xp.arange(10, dtype=xp.bool_)

    @testing.numpy_cupy_array_equal()
    def test_arange_no_dtype_int(self, xp):
        return xp.arange(1, 11, 2)

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_arange_no_dtype_float(self, xp):
        return xp.arange(1.0, 11.0, 2.0)

    @testing.numpy_cupy_array_equal()
    def test_arange_negative_size(self, xp):
        return xp.arange(3, 1)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_linspace(self, xp, dtype):
        return xp.linspace(0, 10, 5, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_linspace2(self, xp, dtype):
        return xp.linspace(10, 0, 5, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    @skip_int_equality_before_numpy_1_20()
    def test_linspace3(self, xp, dtype):
        if xp.dtype(dtype).kind == "u":
            pytest.skip()
        return xp.linspace(-10, 8, 9, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_linspace_zero_num(self, xp, dtype):
        return xp.linspace(0, 10, 0, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_linspace_zero_num_no_endopoint_with_retstep(self, xp, dtype):
        x, step = xp.linspace(
            0, 10, 0, dtype=dtype, endpoint=False, retstep=True
        )
        self.assertTrue(math.isnan(step))
        return x

    @testing.with_requires("numpy>=1.18")
    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_linspace_one_num_no_endopoint_with_retstep(self, xp, dtype):
        start, stop = 3, 7
        x, step = xp.linspace(
            start, stop, 1, dtype=dtype, endpoint=False, retstep=True
        )
        self.assertEqual(step, stop - start)
        return x

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_linspace_one_num(self, xp, dtype):
        return xp.linspace(0, 2, 1, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose()
    def test_linspace_no_endpoint(self, xp, dtype):
        return xp.linspace(0, 10, 5, dtype=dtype, endpoint=False)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_array_equal()
    def test_linspace_with_retstep(self, xp, dtype):
        x, step = xp.linspace(0, 10, 5, dtype=dtype, retstep=True)
        self.assertEqual(step, 2.5)
        return x

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_linspace_no_dtype_int(self, xp):
        return xp.linspace(0, 10, 50)

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_linspace_no_dtype_float(self, xp):
        return xp.linspace(0.0, 10.0, 50)

    @testing.numpy_cupy_array_equal()
    def test_linspace_float_args_with_int_dtype(self, xp):
        return xp.linspace(0.1, 9.1, 11, dtype=int)

    def test_linspace_neg_num(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.linspace(0, 10, -1)

    @testing.numpy_cupy_allclose()
    def test_linspace_float_overflow(self, xp):
        dtype = cupy.default_float_type()
        return xp.linspace(0.0, xp.finfo(dtype).max / 5, 10, dtype=dtype)

    @testing.numpy_cupy_allclose(rtol={numpy.float32: 1e-6, "default": 1e-7})
    def test_linspace_float_underflow(self, xp):
        # find minimum subnormal number
        dtype = cupy.default_float_type()
        # use .tiny instead of .min and while to get
        # minimum subnormal number directly and avoid RuntimeWarning
        x = xp.finfo(dtype).tiny
        return xp.linspace(0.0, x, 10, dtype=dtype)

    @testing.with_requires("numpy>=1.16")
    @testing.for_all_dtypes_combination(
        names=("dtype_range", "dtype_out"), no_bool=True, no_complex=True
    )
    @testing.numpy_cupy_allclose(rtol=1e-04)
    def test_linspace_array_start_stop(self, xp, dtype_range, dtype_out):
        start = xp.array([0, 120], dtype=dtype_range)
        stop = xp.array([100, 0], dtype=dtype_range)
        return xp.linspace(start, stop, num=50, dtype=dtype_out)

    @testing.with_requires("numpy>=1.16")
    @testing.for_all_dtypes_combination(
        names=("dtype_range", "dtype_out"), no_bool=True, no_complex=True
    )
    @testing.numpy_cupy_allclose(rtol=1e-04)
    def test_linspace_mixed_start_stop(self, xp, dtype_range, dtype_out):
        start = 0.0
        if xp.dtype(dtype_range).kind in "u":
            stop = xp.array([100, 16], dtype=dtype_range)
        else:
            stop = xp.array([100, -100], dtype=dtype_range)
        return xp.linspace(start, stop, num=50, dtype=dtype_out)

    @testing.with_requires("numpy>=1.16")
    @testing.for_all_dtypes_combination(
        names=("dtype_range", "dtype_out"), no_bool=True, no_complex=True
    )
    @testing.numpy_cupy_allclose(rtol=1e-04)
    def test_linspace_mixed_start_stop2(self, xp, dtype_range, dtype_out):
        if xp.dtype(dtype_range).kind in "u":
            start = xp.array([160, 120], dtype=dtype_range)
        else:
            start = xp.array([-120, 120], dtype=dtype_range)
        stop = 0
        return xp.linspace(start, stop, num=50, dtype=dtype_out)

    @testing.with_requires("numpy>=1.16")
    @testing.for_all_dtypes_combination(
        names=("dtype_range", "dtype_out"), no_bool=True, no_complex=True
    )
    @testing.numpy_cupy_allclose(rtol=1e-04)
    def test_linspace_array_start_stop_axis1(self, xp, dtype_range, dtype_out):
        start = xp.array([0, 120], dtype=dtype_range)
        stop = xp.array([100, 0], dtype=dtype_range)
        return xp.linspace(start, stop, num=50, dtype=dtype_out, axis=1)

    @testing.with_requires("numpy>=1.16")
    @testing.for_complex_dtypes()
    @testing.numpy_cupy_allclose()
    def test_linspace_complex_start_stop(self, xp, dtype):
        start = xp.array([0, 120], dtype=dtype)
        stop = xp.array([100, 0], dtype=dtype)
        return xp.linspace(start, stop, num=50, dtype=dtype)

    @testing.with_requires("numpy>=1.16")
    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose(rtol=1e-04)
    def test_linspace_start_stop_list(self, xp, dtype):
        start = [0, 0]
        stop = [100, 16]
        return xp.linspace(start, stop, num=50, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose()
    def test_logspace(self, xp, dtype):
        return xp.logspace(0, 2, 5, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose()
    def test_logspace2(self, xp, dtype):
        return xp.logspace(2, 0, 5, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose()
    def test_logspace_zero_num(self, xp, dtype):
        return xp.logspace(0, 2, 0, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose()
    def test_logspace_one_num(self, xp, dtype):
        return xp.logspace(0, 2, 1, dtype=dtype)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose()
    def test_logspace_no_endpoint(self, xp, dtype):
        return xp.logspace(0, 2, 5, dtype=dtype, endpoint=False)

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_logspace_no_dtype_int(self, xp):
        return xp.logspace(0, 2)

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_logspace_no_dtype_float(self, xp):
        return xp.logspace(0.0, 2.0)

    @testing.numpy_cupy_allclose()
    def test_logspace_float_args_with_int_dtype(self, xp):
        return xp.logspace(0.1, 2.1, 11, dtype=int)

    def test_logspace_neg_num(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.logspace(0, 10, -1)

    @testing.for_all_dtypes(no_bool=True)
    @testing.numpy_cupy_allclose(rtol=1e-04)
    def test_logspace_base(self, xp, dtype):
        return xp.logspace(0, 2, 5, base=2.0, dtype=dtype)

    # See #7946 and https://github.com/numpy/numpy/issues/24957
    @testing.with_requires("numpy>=1.16, !=1.25.*, !=1.26.*")
    @testing.for_all_dtypes_combination(
        names=("dtype_range", "dtype_out"), no_bool=True, no_complex=True
    )
    @testing.numpy_cupy_allclose(rtol=1e-6, contiguous_check=False)
    def test_logspace_array_start_stop_axis1(self, xp, dtype_range, dtype_out):
        start = xp.array([0, 2], dtype=dtype_range)
        stop = xp.array([2, 0], dtype=dtype_range)
        return xp.logspace(start, stop, num=5, dtype=dtype_out, axis=1)


@testing.parameterize(
    *testing.product(
        {
            "indexing": ["xy", "ij"],
            "sparse": [False, True],
            "copy": [False, True],
        }
    )
)
class TestMeshgrid(unittest.TestCase):
    @testing.for_all_dtypes()
    def test_meshgrid0(self, dtype):
        out = cupy.meshgrid(
            indexing=self.indexing, sparse=self.sparse, copy=self.copy
        )
        assert out == []

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_meshgrid1(self, xp, dtype):
        x = xp.arange(2).astype(dtype)
        return xp.meshgrid(
            x, indexing=self.indexing, sparse=self.sparse, copy=self.copy
        )

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_meshgrid2(self, xp, dtype):
        x = xp.arange(2).astype(dtype)
        y = xp.arange(3).astype(dtype)
        return xp.meshgrid(
            x, y, indexing=self.indexing, sparse=self.sparse, copy=self.copy
        )

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_meshgrid3(self, xp, dtype):
        x = xp.arange(2).astype(dtype)
        y = xp.arange(3).astype(dtype)
        z = xp.arange(4).astype(dtype)
        return xp.meshgrid(
            x, y, z, indexing=self.indexing, sparse=self.sparse, copy=self.copy
        )


class TestMgrid(unittest.TestCase):
    @testing.numpy_cupy_array_equal()
    def test_mgrid0(self, xp):
        return xp.mgrid[0:]

    @testing.numpy_cupy_array_equal()
    def test_mgrid1(self, xp):
        return xp.mgrid[-10:10]

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_mgrid2(self, xp):
        return xp.mgrid[-10:10:10j]

    @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]
        return xp.mgrid[x:y:10j]

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_mgrid4(self, xp):
        # check len(keys) > 1
        return xp.mgrid[-10:10:10j, -10:10:10j]

    @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]
        y = xp.ones(10)[:, None]
        return xp.mgrid[x:y:10j, x:y:10j]


class TestOgrid(unittest.TestCase):
    @testing.numpy_cupy_array_equal()
    def test_ogrid0(self, xp):
        return xp.ogrid[0:]

    @testing.numpy_cupy_array_equal()
    def test_ogrid1(self, xp):
        return xp.ogrid[-10:10]

    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_ogrid2(self, xp):
        return xp.ogrid[-10:10:10j]

    @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_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_allclose(rtol=1e-4, type_check=has_support_aspect64())
    def test_ogrid5(self, xp):
        # check len(keys) > 1
        x = xp.zeros(10)[:, None]
        y = xp.ones(10)[:, None]
        return xp.ogrid[x:y:10j, x:y:10j]