improve the qubit_states function

Support a more general way of definition:
- use a list or a string for zero/one/plus/minus state
- use a list of integers for defining arbitrary disentangled quantum states (up to a global phase).

src/qutip_qip/qubits.py

@@ -35,37 +35,93 @@
     'qubit_states', 'truncate_to_qubit_state', 'expand_qubit_state']
 
 import qutip
-from qutip import (tensor, basis)
 import numpy as np
 from numpy import sqrt
 
 
-def qubit_states(N=1, states=[0]):
+def qubit_states(states):
     """
-    Function to define initial state of the qubits.
+    Shortcut to generate disentangled qubit states.
 
     Parameters
     ----------
-    N : Integer
-        Number of qubits in the register.
-    states : List
-        Initial state of each qubit.
+    states : list or str
+        - If a list consisting of ``0``, ``1``, ``"0"``, ``"1"``, ``"+"``
+          and ``"-"``,  return the corresponding zero/one/plus/minus state.
+        - If a string consisting of ``0``, ``1``, ``+``, ``-``,
+          same as above.
+        - If a list of float or complex numbers,
+          each number is mapped to a state of the form
+          :math:`\\sqrt{1 - |a|^2} \\left|0\\right\\rangle + a |1\\rangle`,
+          where :math:`a` is the given number.
 
     Returns
-    ----------
-    qstates : Qobj
-        List of qubits.
+    -------
+    quantum_states : :obj:`qutip.Qobj`
+        The generated qubit states.
 
+    Examples
+    --------
+    >>> from qutip_qip.qubits import qubit_states
+    >>> qubit_states([0, 0])  # doctest: +NORMALIZE_WHITESPACE
+    Quantum object: dims = [[2, 2], [1, 1]], shape = (4, 1), type = ket
+    Qobj data =
+    [[1.]
+     [0.]
+     [0.]
+     [0.]]
+    >>> qubit_states([1, "+"])  # doctest: +NORMALIZE_WHITESPACE
+    Quantum object: dims = [[2, 2], [1, 1]], shape =
+    (4, 1), type = ket
+    Qobj data =
+    [[0.        ]
+     [0.        ]
+     [0.70710678]
+     [0.70710678]]
+    >>> qubit_states("-")  # doctest: +NORMALIZE_WHITESPACE
+    Quantum object: dims = [[2], [1]], shape = (2, 1), type = ket
+    Qobj data =
+    [[ 0.70710678]
+     [-0.70710678]]
+    >>> qubit_states("1-")  # doctest: +NORMALIZE_WHITESPACE
+    Quantum object: dims = [[2, 2], [1, 1]], shape =
+    (4, 1), type = ket
+    Qobj data =
+    [[ 0.        ]
+     [ 0.        ]
+     [ 0.70710678]
+     [-0.70710678]]
+    >>> import numpy as np
+    >>> qubit_states([1.j/np.sqrt(2)])  # doctest: +NORMALIZE_WHITESPACE
+    Quantum object: dims = [[2], [1]], shape = (2, 1), type = ket
+    Qobj data =
+    [[0.70710678+0.j        ]
+     [0.        +0.70710678j]]
     """
-    state_list = []
-    for i in range(N):
-        if N > len(states) and i >= len(states):
-            state_list.append(0)
-        else:
-            state_list.append(states[i])
+    if all([np.issubdtype(type(s), np.integer) for s in states]):
+        return qutip.basis([2] * len(states), states)
+
+    states_map = {
+        0: qutip.basis(2, 0),
+        1: qutip.basis(2, 1),
+        "0": qutip.basis(2, 0),
+        "1": qutip.basis(2, 1),
+        "+": (qutip.basis(2, 0) + qutip.basis(2, 1)).unit(),
+        "-": (qutip.basis(2, 0) - qutip.basis(2, 1)).unit(),
+    }
 
-    return tensor([alpha * basis(2, 1) + sqrt(1 - alpha**2) * basis(2, 0)
-                  for alpha in state_list])
+    states_list = []
+    for s in states:
+        if s in states_map:
+            states_list.append(states_map[s])
+        elif np.isscalar(s) and abs(s) <= 1:
+            states_list.append(
+                s * qutip.basis(2, 1) +
+                np.sqrt(1 - abs(s)**2) * qutip.basis(2, 0)
+            )
+        else:
+            raise TypeError(f"Invalid input {s}.")
+    return qutip.tensor(states_list)
 
 
 def _find_reduced_indices(dims):

tests/test_optpulseprocessor.py

@@ -66,8 +66,8 @@ def test_simple_hadamard(self):
             qc, num_tslots=num_tslots, evo_time=evo_time, verbose=True)
 
         # test run_state
-        rho0 = qubit_states(1, [0])
-        plus = (qubit_states(1, [0]) + qubit_states(1, [1])).unit()
+        rho0 = qubit_states([0])
+        plus = (qubit_states([0]) + qubit_states([1])).unit()
         result = test.run_state(rho0)
         assert_allclose(fidelity(result.states[-1], plus), 1, rtol=1.0e-6)
 
@@ -112,8 +112,8 @@ def test_multi_qubits(self):
         qc = [tensor([identity(2), cnot()])]
         test.load_circuit(qc, num_tslots=num_tslots,
                           evo_time=evo_time, min_fid_err=1.0e-6)
-        rho0 = qubit_states(3, [1, 1, 1])
-        rho1 = qubit_states(3, [1, 1, 0])
+        rho0 = qubit_states([1, 1, 1])
+        rho1 = qubit_states([1, 1, 0])
         result = test.run_state(
             rho0, options=Options(store_states=True))
         assert_(fidelity(result.states[-1], rho1) > 1-1.0e-6)

tests/test_processor.py

@@ -278,7 +278,7 @@ def testNoise(self):
         Test for Processor with noise
         """
         # setup and fidelity without noise
-        init_state = qubit_states(2, [0, 0, 0, 0])
+        init_state = qubit_states([0, 0])
         tlist = np.array([0., np.pi/2.])
         a = destroy(2)
         proc = Processor(N=2)
@@ -287,15 +287,15 @@ def testNoise(self):
         proc.pulses[0].coeff = np.array([1])
         result = proc.run_state(init_state=init_state)
         assert_allclose(
-            fidelity(result.states[-1], qubit_states(2, [0, 1, 0, 0])),
+            fidelity(result.states[-1], qubit_states([0, 1])),
             1, rtol=1.e-7)
 
         # decoherence noise
         dec_noise = DecoherenceNoise([0.25*a], targets=1)
         proc.add_noise(dec_noise)
         result = proc.run_state(init_state=init_state)
         assert_allclose(
-            fidelity(result.states[-1], qubit_states(2, [0, 1, 0, 0])),
+            fidelity(result.states[-1], qubit_states([0, 1])),
             0.981852, rtol=1.e-3)
 
         # white random noise
@@ -336,7 +336,7 @@ def testDrift(self):
 
     def testChooseSolver(self):
         # setup and fidelity without noise
-        init_state = qubit_states(2, [0, 0, 0, 0])
+        init_state = qubit_states([0, 0])
         tlist = np.array([0., np.pi/2.])
         a = destroy(2)
         proc = Processor(N=2)
@@ -345,7 +345,7 @@ def testChooseSolver(self):
         proc.pulses[0].coeff = np.array([1])
         result = proc.run_state(init_state=init_state, solver="mcsolve")
         assert_allclose(
-            fidelity(result.states[-1], qubit_states(2, [0, 1, 0, 0])),
+            fidelity(result.states[-1], qubit_states([0, 1])),
             1, rtol=1.e-7)
 
     def test_no_saving_intermidiate_state(self):

tests/test_qubits.py

@@ -30,11 +30,10 @@
 #    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ###############################################################################
-from numpy.testing import assert_, run_module_suite
 import numpy as np
 import pytest
 import qutip
-from qutip import (tensor,  basis) 
+from qutip import tensor, basis
 from qutip_qip.qubits import (
     qubit_states, truncate_to_qubit_state, expand_qubit_state)
 
@@ -47,17 +46,18 @@ def testQubitStates(self):
         """
         Tests the qubit_states function.
         """
-        psi0_a = basis(2, 0)
-        psi0_b = qubit_states()
-        assert_(psi0_a == psi0_b)
-
-        psi1_a = basis(2, 1)
-        psi1_b = qubit_states(states=[1])
-        assert_(psi1_a == psi1_b)
-
-        psi01_a = tensor(psi0_a, psi1_a)
-        psi01_b = qubit_states(N=2, states=[0, 1])
-        assert_(psi01_a == psi01_b)
+        assert(qubit_states([0]) == basis(2, 0))
+        assert(qubit_states([1]) == basis(2, 1))
+        assert(qubit_states([0, 1]) == tensor(basis(2, 0), basis(2, 1)))
+        plus = (basis(2, 0) + basis(2, 1)).unit()
+        minus = (basis(2, 0) - basis(2, 1)).unit()
+        assert(qubit_states("-+") == tensor(minus, plus))
+        assert(qubit_states("0+") == tensor(basis(2, 0), plus))
+        assert(qubit_states("+11") == tensor(plus, basis(2, 1), basis(2, 1)))
+        assert(
+            qubit_states([1.j/np.sqrt(2), 1.]) ==
+            tensor(qutip.Qobj([[1/np.sqrt(2)], [1.j/np.sqrt(2)]]), basis(2, 1))
+            )
 
     @pytest.mark.parametrize(
         "state, full_dims",