Repository: qutip/qutip

Issue:
tensor_contract fails for adjacent dimensions
This one seems to be due to a quirk of NumPy's fancy indexing, but `tensor_contract` fails to contract out a pair `(i, j)` if `j == i + 1 and i > 0`. For example::

``` python
In [1]: import numpy as np

In [2]: import qutip as qt

In [3]: dat = np.arange(2 * 3 * 3 * 4).reshape((2 * 3, 3 * 4))

In [4]: qobj = qt.Qobj(dat, dims=[[2, 3], [3, 4]])

In [5]: qt.tensor_contract(qobj, (1, 2))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-cfa6198a0fa3> in <module>()
----> 1 qt.tensor_contract(qobj, (1, 2))

C:\Anaconda\lib\site-packages\qutip-3.2.0.dev_8af111f-py2.7-win-amd64.egg\qutip\
tensor.pyc in tensor_contract(qobj, *pairs)
    421
    422     # Reshape back into a 2D matrix.
--> 423     qmtx = qtens.reshape((l_mtx_dims, r_mtx_dims))
    424
    425     # Return back as a qobj.

ValueError: total size of new array must be unchanged
```

tensor_contract fails for adjacent dimensions
This one seems to be due to a quirk of NumPy's fancy indexing, but `tensor_contract` fails to contract out a pair `(i, j)` if `j == i + 1 and i > 0`. For example::

``` python
In [1]: import numpy as np

In [2]: import qutip as qt

In [3]: dat = np.arange(2 * 3 * 3 * 4).reshape((2 * 3, 3 * 4))

In [4]: qobj = qt.Qobj(dat, dims=[[2, 3], [3, 4]])

In [5]: qt.tensor_contract(qobj, (1, 2))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-cfa6198a0fa3> in <module>()
----> 1 qt.tensor_contract(qobj, (1, 2))

C:\Anaconda\lib\site-packages\qutip-3.2.0.dev_8af111f-py2.7-win-amd64.egg\qutip\
tensor.pyc in tensor_contract(qobj, *pairs)
    421
    422     # Reshape back into a 2D matrix.
--> 423     qmtx = qtens.reshape((l_mtx_dims, r_mtx_dims))
    424
    425     # Return back as a qobj.

ValueError: total size of new array must be unchanged
```



Discussion:

Resolving patch:
diff --git a/qutip/tensor.py b/qutip/tensor.py
--- a/qutip/tensor.py
+++ b/qutip/tensor.py
@@ -359,7 +359,8 @@ def _tensor_contract_single(arr, i, j):
     idxs = np.arange(arr.shape[i])
     sl = tuple(slice(None, None, None)
                if idx not in (i, j) else idxs for idx in range(arr.ndim))
-    return np.sum(arr[sl], axis=0)
+    contract_at = i if j == i + 1 else 0
+    return np.sum(arr[sl], axis=contract_at)
 
 
 def _tensor_contract_dense(arr, *pairs):


Test patch:
diff --git a/qutip/tests/test_tensor.py b/qutip/tests/test_tensor.py
--- a/qutip/tests/test_tensor.py
+++ b/qutip/tests/test_tensor.py
@@ -31,8 +31,11 @@
 #    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ###############################################################################
 
+import numpy as np
+
 from numpy.testing import assert_equal, assert_, run_module_suite
 
+from qutip.qobj import Qobj
 from qutip.operators import identity
 from qutip.superop_reps import to_super
 from qutip.tensor import (
@@ -68,7 +71,7 @@ def test_unflatten():
     assert unflatten(flatten(l), labels) == l
 
 
-def test_tensor_contract():
+def test_tensor_contract_ident():
     qobj = identity([2, 3, 4])
     ans = 3 * identity([2, 4])
 
@@ -80,5 +83,30 @@ def test_tensor_contract():
     correct_dims = [[[2, 4], [2, 4]], [[2, 4], [2, 4]]]
     assert_equal(correct_dims, tensor_contract(sqobj, (1, 4), (7, 10)).dims)
 
+def case_tensor_contract_other(left, right, i, j, expected_dims, expected_data):
+    dat = np.arange(np.product(left) * np.product(right)).reshape((np.product(left), np.product(right)))
+    
+    qobj = Qobj(dat, dims=[left, right])
+    cqobj = tensor_contract(qobj, (i, j))
+
+    assert_equal(cqobj.dims, expected_dims)
+    assert_equal(cqobj.data.toarray(), expected_data)
+
+def test_tensor_contract_other():
+    yield (
+        case_tensor_contract_other, [2, 3], [3, 4], 1, 2, [[2], [4]],
+        np.einsum('abbc', np.arange(2 * 3 * 3 * 4).reshape((2, 3, 3, 4)))
+    )
+
+    yield (
+        case_tensor_contract_other, [2, 3], [4, 3], 1, 3, [[2], [4]],
+        np.einsum('abcb', np.arange(2 * 3 * 3 * 4).reshape((2, 3, 4, 3)))
+    )
+
+    yield (
+        case_tensor_contract_other, [2, 3], [4, 3], 1, 3, [[2], [4]],
+        np.einsum('abcb', np.arange(2 * 3 * 3 * 4).reshape((2, 3, 4, 3)))
+    )
+
 if __name__ == "__main__":
     run_module_suite()
\ No newline at end of file