Improve initial_state handling for CliffordTableau and StabilizerStat…

…eChForm (#3309)

* Use big_endian_int_to_digits method in CliffordTableau and StabilizerStateChForm and get bounds validation for free

* Update cirq/sim/clifford/clifford_simulator.py

Co-authored-by: Balint Pato <balopat@users.noreply.github.com>

Co-authored-by: Balint Pato <balopat@users.noreply.github.com>

cirq/sim/clifford/clifford_simulator.py

@@ -82,7 +82,8 @@ def _base_iterator(self, circuit: circuits.Circuit,
             qubit_order: Determines the canonical ordering of the qubits. This
                 is often used in specifying the initial state, i.e. the
                 ordering of the computational basis states.
-            initial_state: The initial state for the simulation.
+            initial_state: The initial state for the simulation in the
+                computational basis. Represented as a big endian int.
 
 
         Yields:

cirq/sim/clifford/clifford_simulator_test.py

@@ -179,6 +179,14 @@ def test_simulate_moment_steps_intermediate_measurement():
             np.testing.assert_almost_equal(step.state.to_numpy(), expected)
 
 
+def test_clifford_state_initial_state():
+    q0 = cirq.LineQubit(0)
+    with pytest.raises(ValueError, match='Out of range'):
+        _ = cirq.CliffordState(qubit_map={q0: 0}, initial_state=2)
+    state = cirq.CliffordState(qubit_map={q0: 0}, initial_state=1)
+    np.testing.assert_allclose(state.state_vector(), [0, 1])
+
+
 def test_clifford_trial_result_repr():
     q0 = cirq.LineQubit(0)
     final_simulator_state = cirq.CliffordState(qubit_map={q0: 0})

cirq/sim/clifford/clifford_tableau.py

@@ -18,6 +18,7 @@
 import cirq
 from cirq import protocols
 from cirq.ops.dense_pauli_string import DensePauliString
+from cirq.value import big_endian_int_to_digits
 
 
 class CliffordTableau():
@@ -31,18 +32,22 @@ class CliffordTableau():
     an eigenoperator of the state vector with eigenvalue one: P|psi> = |psi>.
     """
 
-    def __init__(self, num_qubits, initial_state=0):
+    def __init__(self, num_qubits, initial_state: int = 0):
+        """Initializes CliffordTableau
+        Args:
+            num_qubits: The number of qubits in the system.
+            initial_state: The computational basis representation of the
+                state as a big endian int.
+            """
         self.n = num_qubits
 
         self.rs = np.zeros(2 * self.n + 1, dtype=bool)
 
-        def bits(s):
-            while s > 0:
-                yield s & 1
-                s >>= 1
-
-        for (i, val) in enumerate(bits(initial_state)):
-            self.rs[2 * self.n - i - 1] = bool(val)
+        for (i, val) in enumerate(
+                big_endian_int_to_digits(initial_state,
+                                         digit_count=num_qubits,
+                                         base=2)):
+            self.rs[self.n + i] = bool(val)
 
         self.xs = np.zeros((2 * self.n + 1, self.n), dtype=bool)
         self.zs = np.zeros((2 * self.n + 1, self.n), dtype=bool)

cirq/sim/clifford/stabilizer_state_ch_form.py

@@ -17,6 +17,7 @@
 
 import cirq
 from cirq import protocols, value
+from cirq.value import big_endian_int_to_digits
 from cirq._compat import deprecated
 
 
@@ -31,15 +32,12 @@ class StabilizerStateChForm():
     Reference: https://arxiv.org/abs/1808.00128
     """
 
-    def __init__(self,
-                 num_qubits: int,
-                 initial_state: Union[int, np.ndarray] = 0) -> None:
+    def __init__(self, num_qubits: int, initial_state: int = 0) -> None:
         """Initializes StabilizerStateChForm
         Args:
-            num_qubits: The number of qubits in the system
-            initial_state: If an int, the state is set to the computational
-            basis state corresponding to this state.
-            If an np.ndarray it is the full initial state.
+            num_qubits: The number of qubits in the system.
+            initial_state: The computational basis representation of the
+                state as a big endian int.
             """
         self.n = num_qubits
 
@@ -56,15 +54,13 @@ def __init__(self,
 
         self.omega = 1
 
-        def bits(s):
-            while s > 0:
-                yield s & 1
-                s >>= 1
-
         # Apply X for every non-zero element of initial_state
-        for (i, val) in enumerate(bits(initial_state)):
+        for (i, val) in enumerate(
+                big_endian_int_to_digits(initial_state,
+                                         digit_count=num_qubits,
+                                         base=2)):
             if val:
-                self._X(self.n - i - 1)
+                self._X(i)
 
     def _json_dict_(self) -> Dict[str, Any]:
         return protocols.obj_to_dict_helper(

cirq/sim/clifford/stabilizer_state_ch_form_test.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
+import pytest
+
 import cirq
 import cirq.testing
 
@@ -24,3 +27,14 @@ def test_deprecated():
                                   'deprecated'):
         _ = cirq.StabilizerStateChForm(initial_state=0,
                                        num_qubits=1).wave_function()
+
+
+def test_initial_state():
+    with pytest.raises(ValueError, match='Out of range'):
+        _ = cirq.StabilizerStateChForm(initial_state=-31, num_qubits=5)
+    with pytest.raises(ValueError, match='Out of range'):
+        _ = cirq.StabilizerStateChForm(initial_state=32, num_qubits=5)
+    state = cirq.StabilizerStateChForm(initial_state=23, num_qubits=5)
+    expected_state_vector = np.zeros(32)
+    expected_state_vector[23] = 1
+    np.testing.assert_allclose(state.state_vector(), expected_state_vector)