Change plot methods according to recommendations in plotting.md

cirq/experiments/cross_entropy_benchmarking.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Set, Tuple, Sequence, Dict, Any, NamedTuple, Union
-from typing import Iterable
+from typing import (Any, Dict, Iterable, List, NamedTuple, Optional, Sequence,
+                    Set, Tuple, Union)
 import numpy as np
 from matplotlib import pyplot as plt
 from cirq import devices, ops, circuits, sim, work
@@ -44,21 +44,29 @@ def data(self) -> Sequence[CrossEntropyPair]:
         """
         return self._data
 
-    def plot(self, **plot_kwargs: Any) -> None:
+    def plot(self, ax: Optional[plt.Axes] = None,
+             **plot_kwargs: Any) -> plt.Axes:
         """Plots the average XEB fidelity vs the number of cycles.
 
         Args:
-            **plot_kwargs: Arguments to be passed to 'matplotlib.pyplot.plot'.
+            ax: the plt.Axes to plot on. If not given, a new figure is created,
+                plotted on, and shown.
+            **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
+        Returns:
+            The plt.Axes containing the plot.
         """
+        show_plot = not ax
+        if not ax:
+            fig, ax = plt.subplots(1, 1, figsize=(8, 8))
         num_cycles = [d.num_cycle for d in self._data]
         fidelities = [d.xeb_fidelity for d in self._data]
-        fig = plt.figure()
-        ax = plt.gca()
         ax.set_ylim([0, 1.1])
-        plt.plot(num_cycles, fidelities, 'ro-', figure=fig, **plot_kwargs)
-        plt.xlabel('Number of Cycles', figure=fig)
-        plt.ylabel('XEB Fidelity', figure=fig)
-        fig.show(warn=False)
+        ax.plot(num_cycles, fidelities, 'ro-', **plot_kwargs)
+        ax.set_xlabel('Number of Cycles')
+        ax.set_ylabel('XEB Fidelity')
+        if show_plot:
+            fig.show()
+        return ax
 
 
 def cross_entropy_benchmarking(

cirq/experiments/cross_entropy_benchmarking_test.py

@@ -14,6 +14,7 @@
 
 import numpy as np
 
+import matplotlib.pyplot as plt
 import cirq
 
 from cirq.experiments import cross_entropy_benchmarking, build_entangling_layers
@@ -78,4 +79,5 @@ def test_cross_entropy_benchmarking():
     assert len(fidelities_3) == 1
 
     # Sanity test that plot runs.
-    results_1.plot()
+    ax = plt.subplot()
+    results_1.plot(ax)

cirq/experiments/qubit_characterizations.py

@@ -14,12 +14,12 @@
 
 import itertools
 
-from typing import Sequence, Tuple, Iterator, Any, NamedTuple, List
+from typing import Any, Iterator, List, NamedTuple, Optional, Sequence, Tuple
 import numpy as np
 import sympy
 
 from matplotlib import pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D  # type: ignore
+from mpl_toolkits.mplot3d import Axes3D  # type: ignore # pylint: disable=unused-import
 from cirq import circuits, devices, ops, protocols, study, work
 
 Cliffords = NamedTuple('Cliffords',
@@ -54,21 +54,29 @@ def data(self) -> Sequence[Tuple[float, float]]:
         return [(angle, prob) for angle, prob in zip(self._rabi_angles,
                                                      self._excited_state_probs)]
 
-    def plot(self, **plot_kwargs: Any) -> None:
+    def plot(self, ax: Optional[plt.Axes] = None,
+             **plot_kwargs: Any) -> plt.Axes:
         """Plots excited state probability vs the Rabi angle (angle of rotation
         around the x-axis).
 
         Args:
-            **plot_kwargs: Arguments to be passed to matplotlib.pyplot.plot.
+            ax: the plt.Axes to plot on. If not given, a new figure is created,
+                plotted on, and shown.
+            **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
+        Returns:
+            The plt.Axes containing the plot.
         """
-        fig = plt.figure()
-        ax = plt.gca()
+        show_plot = not ax
+        if not ax:
+            fig, ax = plt.subplots(1, 1, figsize=(8, 8))
         ax.set_ylim([0, 1])
-        plt.plot(self._rabi_angles, self._excited_state_probs, 'ro-',
-                 figure=fig, **plot_kwargs)
-        plt.xlabel(r"Rabi Angle (Radian)", figure=fig)
-        plt.ylabel('Excited State Probability', figure=fig)
-        fig.show(warn=False)
+        ax.plot(self._rabi_angles, self._excited_state_probs, 'ro-',
+                **plot_kwargs)
+        ax.set_xlabel(r"Rabi Angle (Radian)")
+        ax.set_ylabel('Excited State Probability')
+        if show_plot:
+            fig.show()
+        return ax
 
 
 class RandomizedBenchMarkResult:
@@ -95,22 +103,28 @@ def data(self) -> Sequence[Tuple[int, float]]:
         return [(num, prob) for num, prob in zip(self._num_cfds_seq,
                                                  self._gnd_state_probs)]
 
-    def plot(self, **plot_kwargs: Any) -> None:
+    def plot(self, ax: Optional[plt.Axes] = None,
+             **plot_kwargs: Any) -> plt.Axes:
         """Plots the average ground state probability vs the number of
         Cliffords in the RB study.
 
         Args:
-            **plot_kwargs: Arguments to be passed to matplotlib.pyplot.plot.
+            ax: the plt.Axes to plot on. If not given, a new figure is created,
+                plotted on, and shown.
+            **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.
+        Returns:
+            The plt.Axes containing the plot.
         """
-        fig = plt.figure()
-        ax = plt.gca()
+        show_plot = not ax
+        if not ax:
+            fig, ax = plt.subplots(1, 1, figsize=(8, 8))
         ax.set_ylim([0, 1])
-
-        plt.plot(self._num_cfds_seq, self._gnd_state_probs, 'ro-',
-                 figure=fig, **plot_kwargs)
-        plt.xlabel(r"Number of Cliffords", figure=fig)
-        plt.ylabel('Ground State Probability', figure=fig)
-        fig.show(warn=False)
+        ax.plot(self._num_cfds_seq, self._gnd_state_probs, 'ro-', **plot_kwargs)
+        ax.set_xlabel(r"Number of Cliffords")
+        ax.set_ylabel('Ground State Probability')
+        if show_plot:
+            fig.show()
+        return ax
 
 
 class TomographyResult:
@@ -130,12 +144,55 @@ def data(self) -> np.ndarray:
         """
         return self._density_matrix
 
-    def plot(self) -> None:
+    def plot(self, axes: Optional[List[plt.Axes]] = None,
+             **plot_kwargs: Any) -> List[plt.Axes]:
         """Plots the real and imaginary parts of the density matrix as two
         3D bar plots.
+
+        Args:
+            axes: a list of 2 `plt.Axes` instances. Note that they must be in
+                3d projections. If not given, a new figure is created with 2
+                axes and the plotted figure is shown.
+            plot_kwargs: the optional kwargs passed to bar3d.
+        Returns:
+            the list of `plt.Axes` being plotted on.
+        Raises:
+            ValueError if axes is a list with length != 2.
         """
-        fig = _plot_density_matrix(self._density_matrix)
-        fig.show(warn=False)
+        show_plot = axes is None
+        if axes is None:
+            fig, axes = plt.subplots(1,
+                                     2,
+                                     figsize=(12.0, 5.0),
+                                     subplot_kw={'projection': '3d'})
+        elif len(axes) != 2:
+            raise ValueError('A TomographyResult needs 2 axes to plot.')
+        mat = self._density_matrix
+        a, _ = mat.shape
+        num_qubits = int(np.log2(a))
+        state_labels = [[0, 1]] * num_qubits
+        kets = []
+        for label in itertools.product(*state_labels):
+            kets.append('|' + str(list(label))[1:-1] + '>')
+        mat_re = np.real(mat)
+        mat_im = np.imag(mat)
+        _matrix_bar_plot(mat_re,
+                         r'Real($\rho$)',
+                         axes[0],
+                         kets,
+                         'Density Matrix (Real Part)',
+                         ylim=(-1, 1),
+                         **plot_kwargs)
+        _matrix_bar_plot(mat_im,
+                         r'Imaginary($\rho$)',
+                         axes[1],
+                         kets,
+                         'Density Matrix (Imaginary Part)',
+                         ylim=(-1, 1),
+                         **plot_kwargs)
+        if show_plot:
+            fig.show()
+        return axes
 
 
 def rabi_oscillations(sampler: work.Sampler,
@@ -489,11 +546,11 @@ def _random_two_q_clifford(q_0: devices.GridQubit, q_1: devices.GridQubit,
 
 def _matrix_bar_plot(mat: np.ndarray,
                      z_label: str,
-                     fig: plt.Figure,
-                     plt_position: int,
+                     ax: plt.Axes,
                      kets: Sequence[str] = None,
                      title: str = None,
-                     ylim: Tuple[int, int] = (-1, 1)) -> None:
+                     ylim: Tuple[int, int] = (-1, 1),
+                     **bar3d_kwargs: Any) -> None:
     num_rows, num_cols = mat.shape
     indices = np.meshgrid(range(num_cols), range(num_rows))
     x_indices = np.array(indices[1]).flatten()
@@ -502,49 +559,26 @@ def _matrix_bar_plot(mat: np.ndarray,
 
     dx = np.ones(mat.size) * 0.3
     dy = np.ones(mat.size) * 0.3
-
-    ax1 = fig.add_subplot(plt_position, projection='3d')  # type: Axes3D
-
     dz = mat.flatten()
-    ax1.bar3d(x_indices, y_indices, z_indices, dx, dy, dz, color='#ff0080',
-              alpha=1.0)
-
-    ax1.set_zlabel(z_label)
-    ax1.set_zlim3d(ylim[0], ylim[1])
+    ax.bar3d(x_indices,
+             y_indices,
+             z_indices,
+             dx,
+             dy,
+             dz,
+             color='#ff0080',
+             alpha=1.0,
+             **bar3d_kwargs)
+
+    ax.set_zlabel(z_label)
+    ax.set_zlim3d(ylim[0], ylim[1])
 
     if kets is not None:
         plt.xticks(np.arange(num_cols) + 0.15, kets)
         plt.yticks(np.arange(num_rows) + 0.15, kets)
 
     if title is not None:
-        ax1.set_title(title)
-
-
-def _plot_density_matrix(mat: np.ndarray) -> plt.Figure:
-    a, _ = mat.shape
-    num_qubits = int(np.log2(a))
-    state_labels = [[0, 1]] * num_qubits
-    kets = []
-    for label in itertools.product(*state_labels):
-        kets.append('|' + str(list(label))[1:-1] + '>')
-    mat_re = np.real(mat)
-    mat_im = np.imag(mat)
-    fig = plt.figure(figsize=(12.0, 5.0))
-    _matrix_bar_plot(mat_re,
-                     r'Real($\rho$)',
-                     fig,
-                     121,
-                     kets,
-                     'Density Matrix (Real Part)',
-                     ylim=(-1, 1))
-    _matrix_bar_plot(mat_im,
-                     r'Imaginary($\rho$)',
-                     fig,
-                     122,
-                     kets,
-                     'Density Matrix (Imaginary Part)',
-                     ylim=(-1, 1))
-    return fig
+        ax.set_title(title)
 
 
 def _two_qubit_clifford(q_0: devices.GridQubit, q_1: devices.GridQubit,

cirq/experiments/qubit_characterizations_test.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 import numpy as np
+import pytest
+
+import matplotlib.pyplot as plt
 
 from cirq import GridQubit
 from cirq import circuits, ops, sim
@@ -142,3 +145,16 @@ def test_two_qubit_state_tomography():
     np.testing.assert_almost_equal(act_rho_hh, tar_rho_hh, decimal=1)
     np.testing.assert_almost_equal(act_rho_xy, tar_rho_xy, decimal=1)
     np.testing.assert_almost_equal(act_rho_yx, tar_rho_yx, decimal=1)
+
+
+def test_tomography_plot_raises_for_incorrect_number_of_axes():
+    simulator = sim.Simulator()
+    qubit = GridQubit(0, 0)
+    circuit = circuits.Circuit(ops.X(qubit)**0.5)
+    result = single_qubit_state_tomography(simulator, qubit, circuit, 1000)
+    with pytest.raises(TypeError):  # ax is not a List[plt.Axes]
+        ax = plt.subplot()
+        result.plot(ax)
+    with pytest.raises(ValueError):
+        _, axes = plt.subplots(1, 3)
+        result.plot(axes)

dev_tools/incremental_coverage.py

@@ -43,6 +43,8 @@
     r'except ImportError',
     # Plotting code.
     r'plt\.show\(\)',
+    r'fig(ure)?\.show\(\)',
+    r'=\s*plt.subplots?\(',
 ]
 EXPLICIT_OPT_OUT_COMMENT = '#coverage:ignore'
 

docs/dev/plotting.md

@@ -28,14 +28,14 @@ interactive session. The recommended way to achieve that is illustrated in the
 example below.
 
 ```python
-from typing import Any, List
+from typing import Any, List, Optional
 import matplotlib.pyplot as plt
 
 class Foo:
     ...
-    def plot(self, ax: plt.Axes=None, **plot_kwargs: Any) -> plt.Axes:
+    def plot(self, ax: Optional[plt.Axes]=None, **plot_kwargs: Any) -> plt.Axes:
         show_plot = not ax
-        if show_plot:
+        if not ax:
             fig, ax = plt.subplots(1, 1)  # or your favorite figure setup
         # Call methods of the ax instance like ax.plot to plot on it.
         ...
@@ -78,10 +78,11 @@ not sufficient. The `plot` method of such a class should take an optional
 ```python
 class Foo:
     ...
-    def plot(self, axes: List[plt.Axes]=None, **plot_kwargs: Any) -> List[plt.Axes]:
+    def plot(self, axes: Optional[List[plt.Axes]]=None,
+             **plot_kwargs: Any) -> List[plt.Axes]:
         show_plot = not axes
-        if show_plot:
-            _, axes = plt.subplots(1, 2)  # or your favorite figure setup
+        if not axes:
+            fig, axes = plt.subplots(1, 2)  # or your favorite figure setup
         elif len(axes) != 2:  # your required number of axes
             raise ValueError('your error message')
         # Call methods of the axes[i] objects to plot on it.