Merge pull request #227 from BoxiLi/expand_operator_csr

Use CSR as default for expand_operator

src/qutip_qip/operations/gates.py

@@ -1,4 +1,5 @@
 import numbers
+from packaging.version import parse as parse_version
 from collections.abc import Iterable
 from itertools import product
 from functools import partial, reduce
@@ -1207,7 +1208,7 @@ def _targets_to_list(targets, oper=None, N=None):
 
 
 def expand_operator(
-    oper, N=None, targets=None, dims=None, cyclic_permutation=False
+    oper, N=None, targets=None, dims=None, cyclic_permutation=False, dtype=None
 ):
     """
     Expand an operator to one that acts on a system with desired dimensions.
@@ -1225,13 +1226,16 @@ def expand_operator(
         The indices of subspace that are acted on.
         Permutation can also be realized by changing the orders of the indices.
     N : int
-         Deprecated. Number of qubits. Please use `dims`.
+        Deprecated. Number of qubits. Please use `dims`.
     cyclic_permutation : boolean, optional
         Deprecated.
         Expand for all cyclic permutation of the targets.
         E.g. if ``N=3`` and `oper` is a 2-qubit operator,
         the result will be a list of three operators,
         each acting on qubits 0 and 1, 1 and 2, 2 and 0.
+    dtype : str, optional
+        Data type of the output `Qobj`. Only for qutip version larger than 5.
+
 
     Returns
     -------
@@ -1276,6 +1280,10 @@ def expand_operator(
      [0. 0. 0. 0. 0. 0. 1. 0.]
      [0. 0. 0. 1. 0. 0. 0. 0.]]
     """
+    if parse_version(qutip.__version__) >= parse_version("5.dev"):
+        # If no data type specified, use CSR
+        dtype = dtype or qutip.settings.core["default_dtype"] or qutip.data.CSR
+        oper = oper.to(dtype)
     if N is not None:
         warnings.warn(
             "The function expand_operator has been generalized to "

tests/test_gates.py

@@ -338,6 +338,19 @@ def test_non_qubit_systems(self, dimensions):
             assert test.dims == expected.dims
             np.testing.assert_allclose(test.full(), expected.full())
 
+    @pytest.mark.skipif(
+        not parse_version(qutip.__version__) >= parse_version('5.dev'),
+        reason=
+        "Data type only exist in v5."
+        )
+    def test_dtype(self):
+        expanded_qobj = expand_operator(gates.cnot(), dims=[2, 2, 2]).data
+        assert isinstance(expanded_qobj, qutip.data.CSR)
+        expanded_qobj = expand_operator(
+            gates.cnot(), dims=[2, 2, 2], dtype="dense").data
+        assert isinstance(expanded_qobj, qutip.data.Dense)
+
+
 def test_gates_class():
     if parse_version(qutip.__version__) < parse_version('5.dev'):
         init_state = qutip.rand_ket(8, dims=[[2, 2, 2], [1, 1, 1]])