Repository: qutip/qutip

Issue:
Extra indices in dims break multiplication
In the development of a new feature (PR coming soon), I ran into a problem where adding new indices to `dims` can break `__mul__` in surprising ways. The full details are in [this notebook](http://nbviewer.ipython.org/gist/cgranade/8f7fce10b6deaf06ce01), but in summary, the bug causes multiplication of Qobjs with dimensions of the form `[[n, 1], [n]]` to not multiply correctly with Qobjs of dimensions `[[n], [n]]`.

Extra indices in dims break multiplication
In the development of a new feature (PR coming soon), I ran into a problem where adding new indices to `dims` can break `__mul__` in surprising ways. The full details are in [this notebook](http://nbviewer.ipython.org/gist/cgranade/8f7fce10b6deaf06ce01), but in summary, the bug causes multiplication of Qobjs with dimensions of the form `[[n, 1], [n]]` to not multiply correctly with Qobjs of dimensions `[[n], [n]]`.



Discussion:

Resolving patch:
diff --git a/qutip/qobj.py b/qutip/qobj.py
--- a/qutip/qobj.py
+++ b/qutip/qobj.py
@@ -62,6 +62,13 @@
 from qutip.sparse import (sp_eigs, sp_expm, sp_fro_norm, sp_max_norm,
                           sp_one_norm, sp_L2_norm)
 
+import sys
+if sys.version_info.major >= 3:
+    from itertools import zip_longest
+elif sys.version_info.major < 3:
+    from itertools import izip_longest
+    zip_longest = izip_longest
+
 
 class Qobj(object):
     """A class for representing quantum objects, such as quantum operators
@@ -396,14 +403,25 @@ def __mul__(self, other):
                 out.data = self.data * other.data
                 dims = [self.dims[0], other.dims[1]]
                 out.dims = dims
+
                 if (not isinstance(dims[0][0], list) and
                         not isinstance(dims[1][0], list)):
-                    r = range(len(dims[0]))
-                    mask = [dims[0][n] == dims[1][n] == 1 for n in r]
-                    out.dims = [max([1], [dims[0][n]
-                                          for n in r if not mask[n]]),
-                                max([1], [dims[1][n]
-                                          for n in r if not mask[n]])]
+                    # If neither left or right is a superoperator,
+                    # we should implicitly partial trace over
+                    # matching dimensions of 1.
+                    # Using izip_longest allows for the left and right dims
+                    # to have uneven length (non-square Qobjs).
+                    # We use None as padding so that it doesn't match anything,
+                    # and will never cause a partial trace on the other side.
+                    mask = [l == r == 1 for l, r in zip_longest(dims[0], dims[1], fillvalue=None)]
+                    # To ensure that there are still any dimensions left, we
+                    # use max() to add a dimensions list of [1] if all matching dims
+                    # are traced out of that side.
+                    out.dims = [max([1], [dim
+                                          for dim, m in zip(dims[0], mask) if not m]),
+                                max([1], [dim
+                                          for dim, m in zip(dims[1], mask) if not m])]
+                    
                 else:
                     out.dims = dims
 


Test patch:
diff --git a/qutip/tests/test_qobj.py b/qutip/tests/test_qobj.py
--- a/qutip/tests/test_qobj.py
+++ b/qutip/tests/test_qobj.py
@@ -174,6 +174,39 @@ def test_QobjDimsShape():
     assert_equal(q1.dims, [[2, 2], [2, 2]])
     assert_equal(q1.shape, [4, 4])
 
+def test_QobjMulNonsquareDims():
+    """
+    Qobj: multiplication w/ non-square qobj.dims
+
+    Checks for regression of #331.
+    """
+    data = np.array([[0, 1], [1, 0]])
+
+    q1 = Qobj(data)
+    q1.dims[0].append(1)
+    q2 = Qobj(data)
+
+    assert_equal((q1 * q2).dims, [[2, 1], [2]])
+    assert_equal((q2 * q1.dag()).dims, [[2], [2, 1]])
+
+    # Note that this is [[2], [2]] instead of [[2, 1], [2, 1]],
+    # as matching dimensions of 1 are implicitly partial traced out.
+    # (See #331.)
+    assert_equal((q1 * q2 * q1.dag()).dims, [[2], [2]])
+    
+    # Because of the above, we also need to check for extra indices
+    # that aren't of length 1.
+    q1 = Qobj([[ 1.+0.j,  0.+0.j],
+         [ 0.+0.j,  1.+0.j],
+         [ 0.+0.j,  1.+0.j],
+         [ 1.+0.j,  0.+0.j],
+         [ 0.+0.j,  0.-1.j],
+         [ 0.+1.j,  0.+0.j],
+         [ 1.+0.j,  0.+0.j],
+         [ 0.+0.j, -1.+0.j]
+     ], dims=[[4, 2], [2]])
+
+    assert_equal((q1 * q2 * q1.dag()).dims, [[4, 2], [4, 2]])
 
 def test_QobjAddition():
     "Qobj addition"