Repository: quantumlib/Cirq

Issue:
StepResult's denisty matrix should not take in integer indices but qubits instead. 
Figuring these out from the natural interfaces of the simulation requires using the qubit ordering or looking into the qubit_map.  The more natural way to do this is to use the qubit_map.


Discussion:
I'll tackle this, I think this is an oversight in one of my earlier PR's, so hopefully a quick fix.

Resolving patch:
diff --git a/cirq/sim/simulator.py b/cirq/sim/simulator.py
--- a/cirq/sim/simulator.py
+++ b/cirq/sim/simulator.py
@@ -781,26 +781,26 @@ def dirac_notation(self, decimals: int = 2) -> str:
             and non-zero floats of the specified accuracy."""
         return wave_function.dirac_notation(self.state(), decimals)
 
-    def density_matrix(self, indices: Iterable[int] = None) -> np.ndarray:
+    def density_matrix(self, qubits: List[ops.QubitId] = None) -> np.ndarray:
         """Returns the density matrix of the wavefunction.
 
-        Calculate the density matrix for the system on the given qubit
-        indices, with the qubits not in indices that are present in self.state
-        traced out. If indices is None the full density matrix for self.state
+        Calculate the density matrix for the system on the list, qubits.
+        Any qubits not in the list that are present in self.state will be
+        traced out. If qubits is None the full density matrix for self.state
         is returned, given self.state follows standard Kronecker convention
         of numpy.kron.
 
         For example:
             self.state = np.array([1/np.sqrt(2), 1/np.sqrt(2)],
                 dtype=np.complex64)
-            indices = None
+            qubits = None
             gives us \rho = \begin{bmatrix}
                                 0.5 & 0.5
                                 0.5 & 0.5
                             \end{bmatrix}
 
         Args:
-            indices: list containing indices for qubits that you would like
+            qubits: list containing qubit IDs that you would like
                 to include in the density matrix (i.e.) qubits that WON'T
                 be traced out.
 
@@ -813,17 +813,19 @@ def density_matrix(self, indices: Iterable[int] = None) -> np.ndarray:
                 corresponding to the state.
         """
         return wave_function.density_matrix_from_state_vector(
-            self.state(), indices)
+            self.state(),
+            [self.qubit_map[q] for q in qubits] if qubits is not None else None
+        )
 
-    def bloch_vector(self, index: int) -> np.ndarray:
+    def bloch_vector(self, qubit: ops.QubitId) -> np.ndarray:
         """Returns the bloch vector of a qubit.
 
-        Calculates the bloch vector of the qubit at index
-        in the wavefunction given by self.state. Given that self.state
+        Calculates the bloch vector of the given qubit
+        in the wavefunction given by self.state, given that self.state
         follows the standard Kronecker convention of numpy.kron.
 
         Args:
-            index: index of qubit who's bloch vector we want to find.
+            qubit: qubit who's bloch vector we want to find.
 
         Returns:
             A length 3 numpy array representing the qubit's bloch vector.
@@ -834,4 +836,4 @@ def bloch_vector(self, index: int) -> np.ndarray:
                 corresponding to the state.
         """
         return wave_function.bloch_vector_from_state_vector(
-            self.state(), index)
+            self.state(), self.qubit_map[qubit])


Test patch:
diff --git a/cirq/sim/simulator_test.py b/cirq/sim/simulator_test.py
--- a/cirq/sim/simulator_test.py
+++ b/cirq/sim/simulator_test.py
@@ -216,53 +216,82 @@ def test_simulator_trial_bloch_vector():
         result.bloch_vector(1))
 
 
-def test_step_result_pretty_state():
-    class BasicStepResult(cirq.StepResult):
+class BasicStepResult(cirq.StepResult):
+
+    def __init__(self, qubit_map: Dict,
+            measurements: Dict[str, List[bool]]) -> None:
+        super().__init__(qubit_map, measurements)
 
-        def __init__(self, qubit_map: Dict,
-                measurements: Dict[str, List[bool]]) -> None:
-            super().__init__(qubit_map, measurements)
+    def state(self) -> np.ndarray:
+        return np.array([0, 1, 0, 0])
 
-        def state(self) -> np.ndarray:
-            return np.array([0, 1, 0, 0])
 
+def test_step_result_pretty_state():
     step_result = BasicStepResult({}, {})
     assert step_result.dirac_notation() == '|01⟩'
 
 
 def test_step_result_density_matrix():
-    class BasicStepResult(cirq.StepResult):
-
-        def __init__(self, qubit_map: Dict,
-                measurements: Dict[str, List[bool]]) -> None:
-            super().__init__(qubit_map, measurements)
+    q0, q1 = cirq.LineQubit.range(2)
 
-        def state(self) -> np.ndarray:
-            return np.array([0, 1, 0, 0])
-
-    step_result = BasicStepResult({}, {})
+    step_result = BasicStepResult({q0: 0, q1: 1}, {})
     rho = np.array([[0, 0, 0, 0],
                     [0, 1, 0, 0],
                     [0, 0, 0, 0],
                     [0, 0, 0, 0]])
+    np.testing.assert_array_almost_equal(rho,
+        step_result.density_matrix([q0, q1]))
+
     np.testing.assert_array_almost_equal(rho,
         step_result.density_matrix())
 
+    rho_ind_rev = np.array([[0, 0, 0, 0],
+                            [0, 0, 0, 0],
+                            [0, 0, 1, 0],
+                            [0, 0, 0, 0]])
+    np.testing.assert_array_almost_equal(rho_ind_rev,
+        step_result.density_matrix([q1, q0]))
 
-def test_step_result_bloch_vector():
-    class BasicStepResult(cirq.StepResult):
+    single_rho = np.array([[0, 0],
+                           [0, 1]])
+    np.testing.assert_array_almost_equal(single_rho,
+        step_result.density_matrix([q1]))
 
-        def __init__(self, qubit_map: Dict,
-                measurements: Dict[str, List[bool]]) -> None:
-            super().__init__(qubit_map, measurements)
 
-        def state(self) -> np.ndarray:
-            return np.array([0, 1, 0, 0])
+def test_step_result_density_matrix_invalid():
+    q0, q1 = cirq.LineQubit.range(2)
 
-    step_result = BasicStepResult({}, {})
-    bloch = np.array([0,0,-1])
-    np.testing.assert_array_almost_equal(bloch,
-        step_result.bloch_vector(1))
+    step_result = BasicStepResult({q0: 0}, {})
+
+    with pytest.raises(KeyError):
+        step_result.density_matrix([q1])
+    with pytest.raises(KeyError):
+        step_result.density_matrix('junk')
+    with pytest.raises(TypeError):
+        step_result.density_matrix(0)
+
+
+def test_step_result_bloch_vector():
+    q0, q1 = cirq.LineQubit.range(2)
+    step_result = BasicStepResult({q0: 0, q1: 1}, {})
+    bloch1 = np.array([0,0,-1])
+    bloch0 = np.array([0,0,1])
+    np.testing.assert_array_almost_equal(bloch1,
+        step_result.bloch_vector(q1))
+    np.testing.assert_array_almost_equal(bloch0,
+        step_result.bloch_vector(q0))
+
+
+def test_bloch_vector_invalid():
+    q0, q1 = cirq.LineQubit.range(2)
+
+    step_result = BasicStepResult({q0: 0}, {})
+    with pytest.raises(KeyError):
+        step_result.bloch_vector(q1)
+    with pytest.raises(KeyError):
+        step_result.bloch_vector('junk')
+    with pytest.raises(KeyError):
+        step_result.bloch_vector(0)
 
 
 class FakeStepResult(cirq.StepResult):