1# This file is part of QuTiP: Quantum Toolbox in Python.
2#
3#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
4#    All rights reserved.
5#
6#    Redistribution and use in source and binary forms, with or without
7#    modification, are permitted provided that the following conditions are
8#    met:
9#
10#    1. Redistributions of source code must retain the above copyright notice,
11#       this list of conditions and the following disclaimer.
12#
13#    2. Redistributions in binary form must reproduce the above copyright
14#       notice, this list of conditions and the following disclaimer in the
15#       documentation and/or other materials provided with the distribution.
16#
17#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
18#       of its contributors may be used to endorse or promote products derived
19#       from this software without specific prior written permission.
20#
21#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
22#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
23#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
24#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
25#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
26#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
27#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
28#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
29#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
30#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
31#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
32###############################################################################
33
34import numpy as np
35
36from numpy.testing import assert_equal, assert_, run_module_suite
37
38from qutip.qobj import Qobj
39from qutip.operators import identity, sigmax, sigmay
40from qutip.superop_reps import to_super, to_choi
41from qutip.random_objects import rand_super_bcsz
42from qutip.tensor import (
43    tensor_contract, tensor_swap
44)
45
46import warnings
47
48def test_tensor_contract_ident():
49    qobj = identity([2, 3, 4])
50    ans = 3 * identity([2, 4])
51
52    assert_(ans == tensor_contract(qobj, (1, 4)))
53
54    # Now try for superoperators.
55    # For now, we just ensure the dims are correct.
56    sqobj = to_super(qobj)
57    correct_dims = [[[2, 4], [2, 4]], [[2, 4], [2, 4]]]
58    assert_equal(correct_dims, tensor_contract(sqobj, (1, 4), (7, 10)).dims)
59
60def case_tensor_contract_other(left, right, pairs, expected_dims, expected_data=None):
61    dat = np.arange(np.product(left) * np.product(right)).reshape((np.product(left), np.product(right)))
62
63    qobj = Qobj(dat, dims=[left, right])
64    cqobj = tensor_contract(qobj, *pairs)
65
66    assert_equal(cqobj.dims, expected_dims)
67    if expected_data is not None:
68        assert_equal(cqobj.data.toarray(), expected_data)
69    else:
70        warnings.warn("tensor_contract test case without checking returned data.")
71
72def test_tensor_contract_other():
73    case_tensor_contract_other([2, 3], [3, 4], [(1, 2)], [[2], [4]],
74        np.einsum('abbc', np.arange(2 * 3 * 3 * 4).reshape((2, 3, 3, 4))))
75
76    case_tensor_contract_other([2, 3], [4, 3], [(1, 3)], [[2], [4]],
77        np.einsum('abcb', np.arange(2 * 3 * 3 * 4).reshape((2, 3, 4, 3))))
78
79    case_tensor_contract_other([2, 3], [4, 3], [(1, 3)], [[2], [4]],
80        np.einsum('abcb', np.arange(2 * 3 * 3 * 4).reshape((2, 3, 4, 3))))
81
82    # Make non-rectangular outputs in a column-/row-symmetric fashion.
83    big_dat = np.arange(2 * 3 * 2 * 3 * 2 * 3 * 2 * 3).reshape((2, 3) * 4)
84
85    case_tensor_contract_other([[2, 3], [2, 3]], [[2, 3], [2, 3]],
86        [(0, 2)], [[[3], [3]], [[2, 3], [2, 3]]],
87        np.einsum('ibidwxyz', big_dat).reshape((3 * 3, 3 * 2 * 3 * 2)))
88
89    case_tensor_contract_other([[2, 3], [2, 3]], [[2, 3], [2, 3]],
90        [(0, 2), (5, 7)], [[[3], [3]], [[2], [2]]],
91        # We separate einsum into two calls due to a bug internal to
92        # einsum.
93        np.einsum('ibidwy', np.einsum('abcdwjyj', big_dat)).reshape((3 * 3, 2 * 2)))
94
95    # Now we try collapsing in a way that's sensitive to column- and row-
96    # stacking conventions.
97    big_dat = np.arange(2 * 2 * 3 * 3 * 2 * 3 * 2 * 3).reshape((3, 3, 2, 2, 2, 3, 2, 3))
98    # Note that the order of [2, 2] and [3, 3] is swapped on the left!
99    big_dims = [[[2, 2], [3, 3]], [[2, 3], [2, 3]]]
100
101    # Let's try a weird tensor contraction; this will likely never come up in practice,
102    # but it should serve as a good canary for more reasonable contractions.
103    case_tensor_contract_other(big_dims[0], big_dims[1],
104        [(0, 4)], [[[2], [3, 3]], [[3], [2, 3]]],
105        # We separate einsum into two calls due to a bug internal to
106        # einsum.
107        np.einsum('abidwxiz', big_dat).reshape((2 * 3 * 3, 3 * 2 * 3)))
108
109
110def case_tensor_swap(qobj, pairs, expected_dims, expected_data=None):
111    sqobj = tensor_swap(qobj, *pairs)
112
113    assert_equal(sqobj.dims, expected_dims)
114    if expected_data is not None:
115        assert_equal(sqobj.data.toarray(), expected_data.data.toarray())
116    else:
117        warnings.warn("tensor_contract test case without checking returned data.")
118
119
120def test_tensor_swap_other():
121    dims = (2, 3, 4, 5, 7)
122
123    for dim in dims:
124        S = to_super(rand_super_bcsz(dim))
125
126        # Swapping the inner indices on a superoperator should give a Choi matrix.
127        J = to_choi(S)
128        case_tensor_swap(S, [(1, 2)], [[[dim], [dim]], [[dim], [dim]]], J)
129
130
131if __name__ == "__main__":
132    run_module_suite()