Repository: qutip/qutip

Issue:
ptrace inconsistent between sparse and dense
`Qobj.ptrace` is inconsistent between `sparse=True` and `sparse=False` arguments if the selection arguments are not in order:

```python
>>> import qutip
>>> a, b, c = qutip.rand_ket(2), qutip.rand_ket(3), qutip.rand_ket(4)
>>> sparse = qutip.tensor(a, b, c).ptrace([2, 1], sparse=True)
>>> dense = qutip.tensor(a, b, c).ptrace([2, 1], sparse=False)
>>> sparse.dims == dense.dims
False
>>> sparse.dims
[[4, 3], [4, 3]]
>>> dense.dims
[[3, 4], [3, 4]]
```

I'm fairly sure the only problem is that the dimensions come out incorrectly - for separable states, we find
```python
>>> ptraced = qutip.tensor(a, b, c).ptrace([2, 1])
>>> built = a.proj().tr() * qutip.tensor(b.proj(), c.proj())
>>> np.allclose(ptraced.full(), built.full())
True
```
which is also consistent with `_ptrace_dense` (which sorts `sel`).

Pretty sure the solution is just to sort `sel` on input, but wanted to check with @Ericgig.


Discussion:
Also, I was just quickly searching for a quick-and-dirty "obviously correct" method of doing the partial trace to use as the test case (rather than keeping around a legacy Cython version just to test against), and found [this StackOverflow answer](https://scicomp.stackexchange.com/questions/27496/calculating-partial-trace-of-array-in-numpy) which has like 3 people telling someone "don't do it yourself, just use QuTiP"!  Pretty neat!
I looked at the documentation and there is no mention about which is expected.

But sparse seems in the wrong here since only the dims are switched, not the data.
I would have expected:
```
ptraced = qutip.tensor(a, b, c).ptrace([2, 1])
built = a.proj().tr() * qutip.tensor(c.proj(), b.proj())
ptraced == built
```
Yes, just sorting `sel` seems appropriate.

Resolving patch:
diff --git a/qutip/qobj.py b/qutip/qobj.py
--- a/qutip/qobj.py
+++ b/qutip/qobj.py
@@ -1329,6 +1329,8 @@ def ptrace(self, sel, sparse=None):
         ----------
         sel : int/list
             An ``int`` or ``list`` of components to keep after partial trace.
+            The order is unimportant; no transposition will be done and the
+            spaces will remain in the same order in the output.
 
         Returns
         -------
@@ -2201,25 +2203,34 @@ def _ptrace_dense(Q, sel):
     else:
         sel = np.asarray(sel)
     sel = list(np.sort(sel))
+    for x in sel:
+        if not 0 <= x < len(rd):
+            raise IndexError("Invalid selection index in ptrace.")
     dkeep = (rd[sel]).tolist()
     qtrace = list(set(np.arange(nd)) - set(sel))
     dtrace = (rd[qtrace]).tolist()
+    if len(dkeep) + len(dtrace) != len(rd):
+        raise ValueError("Duplicate selection index in ptrace.")
+    if not dtrace:
+        # If we are keeping all dimensions, no need to construct an ndarray.
+        return Q.copy()
     rd = list(rd)
     if isket(Q):
         vmat = (Q.full()
                 .reshape(rd)
                 .transpose(sel + qtrace)
-                .reshape([np.prod(dkeep), np.prod(dtrace)]))
+                .reshape([np.prod(dkeep, dtype=np.int32),
+                          np.prod(dtrace, dtype=np.int32)]))
         rhomat = vmat.dot(vmat.conj().T)
     else:
         rhomat = np.trace(Q.full()
                           .reshape(rd + rd)
                           .transpose(qtrace + [nd + q for q in qtrace] +
                                      sel + [nd + q for q in sel])
-                          .reshape([np.prod(dtrace),
-                                    np.prod(dtrace),
-                                    np.prod(dkeep),
-                                    np.prod(dkeep)]))
+                          .reshape([np.prod(dtrace, dtype=np.int32),
+                                    np.prod(dtrace, dtype=np.int32),
+                                    np.prod(dkeep, dtype=np.int32),
+                                    np.prod(dkeep, dtype=np.int32)]))
     return Qobj(rhomat, dims=[dkeep, dkeep])
 
 


Test patch:
diff --git a/qutip/tests/test_ptrace.py b/qutip/tests/test_ptrace.py
--- a/qutip/tests/test_ptrace.py
+++ b/qutip/tests/test_ptrace.py
@@ -30,10 +30,76 @@
 #    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ###############################################################################
+import itertools
+
+import numpy as np
 from numpy.testing import assert_
+import pytest
+
 from qutip import *
 from qutip.legacy.ptrace import _ptrace as _pt
 
+
+@pytest.fixture(params=[True, False], ids=['sparse', 'dense'])
+def sparse(request):
+    return request.param
+
+
+@pytest.fixture(params=[True, False], ids=['dm', 'ket'])
+def dm(request):
+    return request.param
+
+
+@pytest.fixture
+def state(dm):
+    dims = [2, 3, 4]
+    state = qutip.rand_ket(np.prod(dims), dims=[dims, [1]*len(dims)])
+    if dm:
+        state = state.proj()
+    return state
+
+
+def test_ptrace_noncompound_rand(sparse, dm):
+    """Test `A.ptrace(0) == A` when `A` is in a non-tensored Hilbert space."""
+    for _ in range(5):
+        state = qutip.rand_ket(5)
+        if dm:
+            state = state.proj()
+        assert state.ptrace(0, sparse=sparse) == state
+
+
+@pytest.mark.parametrize('pair', list(itertools.combinations(range(3), 2)))
+def test_ptrace_unsorted_selection_subset(state, sparse, pair):
+    """
+    Regression test for gh-1325.  ptrace should work the same independently of
+    the order of the input; no transposition in done in the trace operation.
+    """
+    # pair is always sorted.
+    state_ordered = state.ptrace(pair, sparse=sparse)
+    state_reversed = state.ptrace(pair[::-1], sparse=sparse)
+    assert state_ordered.dims == state_reversed.dims
+    assert state_ordered == state_reversed
+
+
+@pytest.mark.parametrize('permutation', list(itertools.permutations(range(3))))
+def test_ptrace_unsorted_selection_all(state, sparse, permutation):
+    state_ptraced = state.ptrace(permutation, sparse=sparse)
+    assert state.dims == state_ptraced.dims
+    assert state == state_ptraced
+
+
+@pytest.mark.parametrize(['selection', 'exception'], [
+    pytest.param(4, IndexError, id='too big'),
+    pytest.param(-1, IndexError, id='too small'),
+    pytest.param([0, 0], ValueError, id='duplicate'),
+    # 'too many' may throw either from duplication or invalid index.
+    pytest.param([0, 1, 2, 3], Exception, id='too many'),
+])
+def test_ptrace_fails_on_invalid_input(state, sparse, selection, exception):
+    with pytest.raises(exception):
+        state.ptrace(selection, sparse=sparse)
+
+
 def test_ptrace_rand():
     'ptrace : randomized tests, sparse'
     for k in range(5):