<reponame>qutip/qutip-cupy<gh_stars>1-10
import numpy as np
import pytest

from qutip_cupy import CuPyDense
from qutip_cupy import dense

# import qutip after qutip_cupy to supress the matplotlib import warning
from qutip.core import data


@pytest.fixture(scope="function", params=((1, 2), (5, 10), (7, 3), (2, 5)))
def shape(request):
    return request.param


class TestCuPyDenseDispatch:
    """ Tests if the methods and conversions have been
        succesfully registered to QuTiP's Data Layer."""

    def test_conversion_cycle(self, shape):

        qutip_dense = data.Dense(np.random.uniform(size=shape))

        tr1 = data.to(CuPyDense, qutip_dense)
        tr2 = data.to(data.Dense, tr1)

        np.testing.assert_array_equal(qutip_dense.to_array(), tr2.to_array())


class TestCuPyDense:
    """ Tests of the methods and constructors of the CuPyDense class. """

    @pytest.mark.parametrize(["shape"], [pytest.param((5,)), pytest.param((3))])
    def test_init_promotion_to_ket(sef, shape):
        cupy_dense = CuPyDense(np.random.uniform(size=shape))
        assert len(cupy_dense.shape) == 2
        if isinstance(shape, int):
            size = shape
        else:
            size = shape[0]
        assert cupy_dense.shape == (size, 1)

    def test_shape(self, shape):
        cupy_dense = CuPyDense(np.random.uniform(size=shape))
        assert cupy_dense.shape == shape

    def test_transpose(self, shape):
        cupy_dense = CuPyDense(np.random.uniform(size=shape)).transpose()
        np.testing.assert_array_equal(cupy_dense.shape, (shape[1], shape[0]))

    def test_adjoint(self, shape):
        data = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)
        cpd_adj = CuPyDense(data).adjoint()
        np.testing.assert_array_equal(cpd_adj.to_array(), data.transpose().conj())

    @pytest.mark.parametrize(
        ["matrix", "trace"],
        [
            pytest.param([[0, 1], [1, 0]], 0),
            pytest.param([[2.0j, 1], [1, 1]], 1 + 2.0j),
        ],
    )
    def test_trace(self, matrix, trace):
        cupy_array = CuPyDense(matrix)
        assert cupy_array.trace() == trace


def test_true_div(shape):

    array = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)

    cup_arr = CuPyDense(array)
    cpdense_over_2 = cup_arr / 2.0
    qtpdense_over_2 = data.Dense(array) / 2.0

    np.testing.assert_array_equal(cpdense_over_2.to_array(), qtpdense_over_2.to_array())


def test_itrue_div(shape):

    array = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)
    cpd_array = CuPyDense(array)

    cpd_array /= 2
    qtpdense = data.Dense(array / 2)

    np.testing.assert_array_equal(cpd_array.to_array(), qtpdense.to_array())


def test_mul(shape):

    array = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)

    cpdense_mul = CuPyDense(array) * (2.0 + 1.0j)
    qtpdense_mul = data.Dense(array) * (2.0 + 1.0j)

    np.testing.assert_array_equal(cpdense_mul.to_array(), qtpdense_mul.to_array())


def test_matmul(shape):

    array1 = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)

    cpdense1 = CuPyDense(array1)
    qtpdense1 = data.Dense(array1)

    shape2 = (shape[1], np.random.choice([1, 2, 6, 7, 8]))
    array2 = np.random.uniform(size=shape2) + 1.0j * np.random.uniform(size=shape2)

    cpdense2 = CuPyDense(array2)
    qtpdense2 = data.Dense(array2)

    cpdense_matmul = cpdense1 @ cpdense2
    qtpdense_matmul = qtpdense1 @ qtpdense2

    np.testing.assert_array_almost_equal(
        cpdense_matmul.to_array(), qtpdense_matmul.to_array(), decimal=10
    )


class TestFactoryMethods:
    def test_zeros(self, shape):
        base = dense.zeros(shape[0], shape[1])
        nd = base.to_array()
        assert isinstance(base, CuPyDense)
        assert base.shape == shape
        assert nd.shape == shape
        assert np.count_nonzero(nd) == 0

    @pytest.mark.parametrize("dimension", [1, 5, 100])
    @pytest.mark.parametrize(
        "scale",
        [None, 2, -0.1, 1.5, 1.5 + 1j],
        ids=["none", "int", "negative", "float", "complex"],
    )
    def test_identity(self, dimension, scale):
        # scale=None is testing that the default value returns the identity.
        base = (
            dense.identity(dimension)
            if scale is None
            else dense.identity(dimension, scale)
        )
        nd = base.to_array()
        numpy_test = np.eye(dimension, dtype=np.complex128)
        if scale is not None:
            numpy_test *= scale
        assert isinstance(base, CuPyDense)
        assert base.shape == (dimension, dimension)
        assert np.count_nonzero(nd - numpy_test) == 0

    @pytest.mark.parametrize(
        ["diagonals", "offsets", "shape"],
        [
            pytest.param([2j, 3, 5, 9], None, None, id="main diagonal"),
            pytest.param([1], None, None, id="1x1"),
            pytest.param([[0.2j, 0.3]], None, None, id="main diagonal list"),
            pytest.param([0.2j, 0.3], 2, None, id="superdiagonal"),
            pytest.param([0.2j, 0.3], -2, None, id="subdiagonal"),
            pytest.param(
                [[0.2, 0.3, 0.4], [0.1, 0.9]], [-2, 3], None, id="two diagonals"
            ),
            pytest.param([1, 2, 3], 0, (3, 5), id="main wide"),
            pytest.param([1, 2, 3], 0, (5, 3), id="main tall"),
            pytest.param([[1, 2, 3], [4, 5]], [-1, -2], (4, 8), id="two wide sub"),
            pytest.param(
                [[1, 2, 3, 4], [4, 5, 4j, 1j]], [1, 2], (4, 8), id="two wide super"
            ),
            pytest.param([[1, 2, 3], [4, 5]], [1, 2], (8, 4), id="two tall super"),
            pytest.param(
                [[1, 2, 3, 4], [4, 5, 4j, 1j]], [-1, -2], (8, 4), id="two tall sub"
            ),
            pytest.param(
                [[1, 2, 3], [4, 5, 6], [1, 2]], [1, -1, -2], (4, 4), id="out of order"
            ),
            pytest.param(
                [[1, 2, 3], [4, 5, 6], [1, 2]], [1, 1, -2], (4, 4), id="sum duplicates"
            ),
        ],
    )
    def test_diags(self, diagonals, offsets, shape):
        base = dense.diags(diagonals, offsets, shape)
        # Build numpy version test.
        if not isinstance(diagonals[0], list):
            diagonals = [diagonals]
        offsets = np.atleast_1d(offsets if offsets is not None else [0])
        if shape is None:
            size = len(diagonals[0]) + abs(offsets[0])
            shape = (size, size)
        test = np.zeros(shape, dtype=np.complex128)
        for diagonal, offset in zip(diagonals, offsets):
            test[np.where(np.eye(*shape, k=offset) == 1)] += diagonal
        assert isinstance(base, CuPyDense)
        assert base.shape == shape
        np.testing.assert_allclose(base.to_array(), test, rtol=1e-10)

import numpy as np
import pytest

from qutip_cupy import CuPyDense
from qutip_cupy import dense

# import qutip after qutip_cupy to supress the matplotlib import warning
from qutip.core import data


@pytest.fixture(scope="function", params=((1, 2), (5, 10), (7, 3), (2, 5)))
def shape(request):
    return request.param


class TestCuPyDenseDispatch:
    """ Tests if the methods and conversions have been
        succesfully registered to QuTiP's Data Layer."""

    def test_conversion_cycle(self, shape):

        qutip_dense = data.Dense(np.random.uniform(size=shape))

        tr1 = data.to(CuPyDense, qutip_dense)
        tr2 = data.to(data.Dense, tr1)

        np.testing.assert_array_equal(qutip_dense.to_array(), tr2.to_array())


class TestCuPyDense:
    """ Tests of the methods and constructors of the CuPyDense class. """

    @pytest.mark.parametrize(["shape"], [pytest.param((5,)), pytest.param((3))])
    def test_init_promotion_to_ket(sef, shape):
        cupy_dense = CuPyDense(np.random.uniform(size=shape))
        assert len(cupy_dense.shape) == 2
        if isinstance(shape, int):
            size = shape
        else:
            size = shape[0]
        assert cupy_dense.shape == (size, 1)

    def test_shape(self, shape):
        cupy_dense = CuPyDense(np.random.uniform(size=shape))
        assert cupy_dense.shape == shape

    def test_transpose(self, shape):
        cupy_dense = CuPyDense(np.random.uniform(size=shape)).transpose()
        np.testing.assert_array_equal(cupy_dense.shape, (shape[1], shape[0]))

    def test_adjoint(self, shape):
        data = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)
        cpd_adj = CuPyDense(data).adjoint()
        np.testing.assert_array_equal(cpd_adj.to_array(), data.transpose().conj())

    @pytest.mark.parametrize(
        ["matrix", "trace"],
        [
            pytest.param([[0, 1], [1, 0]], 0),
            pytest.param([[2.0j, 1], [1, 1]], 1 + 2.0j),
        ],
    )
    def test_trace(self, matrix, trace):
        cupy_array = CuPyDense(matrix)
        assert cupy_array.trace() == trace


def test_true_div(shape):

    array = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)

    cup_arr = CuPyDense(array)
    cpdense_over_2 = cup_arr / 2.0
    qtpdense_over_2 = data.Dense(array) / 2.0

    np.testing.assert_array_equal(cpdense_over_2.to_array(), qtpdense_over_2.to_array())


def test_itrue_div(shape):

    array = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)
    cpd_array = CuPyDense(array)

    cpd_array /= 2
    qtpdense = data.Dense(array / 2)

    np.testing.assert_array_equal(cpd_array.to_array(), qtpdense.to_array())


def test_mul(shape):

    array = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)

    cpdense_mul = CuPyDense(array) * (2.0 + 1.0j)
    qtpdense_mul = data.Dense(array) * (2.0 + 1.0j)

    np.testing.assert_array_equal(cpdense_mul.to_array(), qtpdense_mul.to_array())


def test_matmul(shape):

    array1 = np.random.uniform(size=shape) + 1.0j * np.random.uniform(size=shape)

    cpdense1 = CuPyDense(array1)
    qtpdense1 = data.Dense(array1)

    shape2 = (shape[1], np.random.choice([1, 2, 6, 7, 8]))
    array2 = np.random.uniform(size=shape2) + 1.0j * np.random.uniform(size=shape2)

    cpdense2 = CuPyDense(array2)
    qtpdense2 = data.Dense(array2)

    cpdense_matmul = cpdense1 @ cpdense2
    qtpdense_matmul = qtpdense1 @ qtpdense2

    np.testing.assert_array_almost_equal(
        cpdense_matmul.to_array(), qtpdense_matmul.to_array(), decimal=10
    )


class TestFactoryMethods:
    def test_zeros(self, shape):
        base = dense.zeros(shape[0], shape[1])
        nd = base.to_array()
        assert isinstance(base, CuPyDense)
        assert base.shape == shape
        assert nd.shape == shape
        assert np.count_nonzero(nd) == 0

    @pytest.mark.parametrize("dimension", [1, 5, 100])
    @pytest.mark.parametrize(
        "scale",
        [None, 2, -0.1, 1.5, 1.5 + 1j],
        ids=["none", "int", "negative", "float", "complex"],
    )
    def test_identity(self, dimension, scale):
        # scale=None is testing that the default value returns the identity.
        base = (
            dense.identity(dimension)
            if scale is None
            else dense.identity(dimension, scale)
        )
        nd = base.to_array()
        numpy_test = np.eye(dimension, dtype=np.complex128)
        if scale is not None:
            numpy_test *= scale
        assert isinstance(base, CuPyDense)
        assert base.shape == (dimension, dimension)
        assert np.count_nonzero(nd - numpy_test) == 0

    @pytest.mark.parametrize(
        ["diagonals", "offsets", "shape"],
        [
            pytest.param([2j, 3, 5, 9], None, None, id="main diagonal"),
            pytest.param([1], None, None, id="1x1"),
            pytest.param([[0.2j, 0.3]], None, None, id="main diagonal list"),
            pytest.param([0.2j, 0.3], 2, None, id="superdiagonal"),
            pytest.param([0.2j, 0.3], -2, None, id="subdiagonal"),
            pytest.param(
                [[0.2, 0.3, 0.4], [0.1, 0.9]], [-2, 3], None, id="two diagonals"
            ),
            pytest.param([1, 2, 3], 0, (3, 5), id="main wide"),
            pytest.param([1, 2, 3], 0, (5, 3), id="main tall"),
            pytest.param([[1, 2, 3], [4, 5]], [-1, -2], (4, 8), id="two wide sub"),
            pytest.param(
                [[1, 2, 3, 4], [4, 5, 4j, 1j]], [1, 2], (4, 8), id="two wide super"
            ),
            pytest.param([[1, 2, 3], [4, 5]], [1, 2], (8, 4), id="two tall super"),
            pytest.param(
                [[1, 2, 3, 4], [4, 5, 4j, 1j]], [-1, -2], (8, 4), id="two tall sub"
            ),
            pytest.param(
                [[1, 2, 3], [4, 5, 6], [1, 2]], [1, -1, -2], (4, 4), id="out of order"
            ),
            pytest.param(
                [[1, 2, 3], [4, 5, 6], [1, 2]], [1, 1, -2], (4, 4), id="sum duplicates"
            ),
        ],
    )
    def test_diags(self, diagonals, offsets, shape):
        base = dense.diags(diagonals, offsets, shape)
        # Build numpy version test.
        if not isinstance(diagonals[0], list):
            diagonals = [diagonals]
        offsets = np.atleast_1d(offsets if offsets is not None else [0])
        if shape is None:
            size = len(diagonals[0]) + abs(offsets[0])
            shape = (size, size)
        test = np.zeros(shape, dtype=np.complex128)
        for diagonal, offset in zip(diagonals, offsets):
            test[np.where(np.eye(*shape, k=offset) == 1)] += diagonal
        assert isinstance(base, CuPyDense)
        assert base.shape == shape
        np.testing.assert_allclose(base.to_array(), test, rtol=1e-10)
# import qutip after qutip_cupy to supress the matplotlib import warning Tests if the methods and conversions have been         succesfully registered to QuTiP's Data Layer. Tests of the methods and constructors of the CuPyDense class. # scale=None is testing that the default value returns the identity. # Build numpy version test.