Repository: quantumlib/Cirq

Issue:
Support multi-qubit measurements in `cirq.plot_state_histogram`
Quote from the docstring: "Currently this function assumes each measurement gate applies to only a single qubit."

Currently, I get an incorrect histogram if I didn't read the docstring and used a multi-qubit measurement (I always make circuits like this):
```python
qubits = cirq.LineQubit.range(3)
c = cirq.Circuit(
    (cirq.X**0.4).on_each(*qubits),
    cirq.measure(*qubits),  # One multi-qubit measurement
)
cirq.plot_state_histogram(cirq.sample(c, repetitions=10000))
# Incorrect output, no warning or error
```
![index2](https://user-images.githubusercontent.com/2476062/82605089-e5c42400-9ba4-11ea-8eb5-516586620a45.png)

If I use single-qubit measurement gates, I get the expected histogram:
```python
qubits = cirq.LineQubit.range(3)
c = cirq.Circuit(
    (cirq.X**0.4).on_each(*qubits),
    cirq.measure_each(*qubits),  # One measurement per qubit
)
cirq.plot_state_histogram(cirq.sample(c, repetitions=10000))
```
![index](https://user-images.githubusercontent.com/2476062/82605108-ebba0500-9ba4-11ea-9752-f337ecab1d26.png)

This looks like it could be fixed by adding some logic to `plot_state_histogram` (https://github.com/quantumlib/Cirq/blob/master/cirq/study/visualize.py#L22) that checks for multi-qubit measurements and either correctly interpret them or raise an error.


Discussion:

Resolving patch:
diff --git a/cirq/study/visualize.py b/cirq/study/visualize.py
--- a/cirq/study/visualize.py
+++ b/cirq/study/visualize.py
@@ -24,8 +24,6 @@ def plot_state_histogram(result: trial_result.TrialResult) -> np.ndarray:
 
     States is a bitstring representation of all the qubit states in a single
     result.
-    Currently this function assumes each measurement gate applies to only
-    a single qubit.
 
     Args:
         result: The trial results to plot.
@@ -39,17 +37,15 @@ def plot_state_histogram(result: trial_result.TrialResult) -> np.ndarray:
     # This allows cirq to be usable without python3-tk.
     import matplotlib.pyplot as plt
 
-    num_qubits = len(result.measurements.keys())
+    num_qubits = sum([value.shape[1] for value in result.measurements.values()])
     states = 2**num_qubits
     values = np.zeros(states)
-
     # measurements is a dict of {measurement gate key:
     #                            array(repetitions, boolean result)}
     # Convert this to an array of repetitions, each with an array of booleans.
     # e.g. {q1: array([[True, True]]), q2: array([[False, False]])}
     #      --> array([[True, False], [True, False]])
-    measurement_by_result = np.array([
-        v.transpose()[0] for k, v in result.measurements.items()]).transpose()
+    measurement_by_result = np.hstack(list(result.measurements.values()))
 
     for meas in measurement_by_result:
         # Convert each array of booleans to a string representation.


Test patch:
diff --git a/cirq/study/visualize_test.py b/cirq/study/visualize_test.py
--- a/cirq/study/visualize_test.py
+++ b/cirq/study/visualize_test.py
@@ -38,3 +38,30 @@ def test_plot_state_histogram():
     expected_values = [0., 0., 0., 5.]
 
     np.testing.assert_equal(values_plotted, expected_values)
+
+
+def test_plot_state_histogram_multi_1():
+    pl.switch_backend('PDF')
+    qubits = cirq.LineQubit.range(4)
+    c = cirq.Circuit(
+        cirq.X.on_each(*qubits[1:]),
+        cirq.measure(*qubits),  # One multi-qubit measurement
+    )
+    r = cirq.sample(c, repetitions=5)
+    values_plotted = visualize.plot_state_histogram(r)
+    expected_values = [0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0]
+    np.testing.assert_equal(values_plotted, expected_values)
+
+
+def test_plot_state_histogram_multi_2():
+    pl.switch_backend('PDF')
+    qubits = cirq.LineQubit.range(4)
+    c = cirq.Circuit(
+        cirq.X.on_each(*qubits[1:]),
+        cirq.measure(*qubits[:2]),  # One multi-qubit measurement
+        cirq.measure_each(*qubits[2:]),  # Multiple single-qubit measurement
+    )
+    r = cirq.sample(c, repetitions=5)
+    values_plotted = visualize.plot_state_histogram(r)
+    expected_values = [0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0]
+    np.testing.assert_equal(values_plotted, expected_values)