From 696bf302dfda85d4c0f3154a7c1133b8e2b7416f Mon Sep 17 00:00:00 2001
From: Ping Yeh <pingyeh@google.com>
Date: Tue, 8 Oct 2019 19:50:55 -0700
Subject: [PATCH 1/7] Add ax arg and remove calling figure.show() in plot
 methods.

---
 .../experiments/cross_entropy_benchmarking.py | 11 +++-----
 .../cross_entropy_benchmarking_test.py        |  5 +++-
 cirq/experiments/qubit_characterizations.py   | 25 +++++++------------
 .../cross_entropy_benchmarking_example.py     |  6 ++++-
 examples/qubit_characterizations_example.py   | 10 +++++---
 5 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/cirq/experiments/cross_entropy_benchmarking.py b/cirq/experiments/cross_entropy_benchmarking.py
index 0893e93e095..4d6031815de 100644
--- a/cirq/experiments/cross_entropy_benchmarking.py
+++ b/cirq/experiments/cross_entropy_benchmarking.py
@@ -44,7 +44,7 @@ def data(self) -> Sequence[CrossEntropyPair]:
         """
         return self._data
 
-    def plot(self, **plot_kwargs: Any) -> None:
+    def plot(self, ax: plt.Axes, **plot_kwargs: Any) -> None:
         """Plots the average XEB fidelity vs the number of cycles.
 
         Args:
@@ -52,13 +52,10 @@ def plot(self, **plot_kwargs: Any) -> None:
         """
         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')
 
 
 def cross_entropy_benchmarking(
diff --git a/cirq/experiments/cross_entropy_benchmarking_test.py b/cirq/experiments/cross_entropy_benchmarking_test.py
index 93b1f9a4e72..46146db8710 100644
--- a/cirq/experiments/cross_entropy_benchmarking_test.py
+++ b/cirq/experiments/cross_entropy_benchmarking_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import matplotlib as mpl
 import numpy as np
 
 import cirq
@@ -78,4 +79,6 @@ def test_cross_entropy_benchmarking():
     assert len(fidelities_3) == 1
 
     # Sanity test that plot runs.
-    results_1.plot()
+    figure = mpl.figure.Figure()
+    ax = figure.add_subplot(111)
+    results_1.plot(ax)
diff --git a/cirq/experiments/qubit_characterizations.py b/cirq/experiments/qubit_characterizations.py
index 146238d2dc9..cdfa0ca197e 100644
--- a/cirq/experiments/qubit_characterizations.py
+++ b/cirq/experiments/qubit_characterizations.py
@@ -54,21 +54,18 @@ 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: plt.Axes, **plot_kwargs: Any) -> None:
         """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.
         """
-        fig = plt.figure()
-        ax = plt.gca()
         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')
 
 
 class RandomizedBenchMarkResult:
@@ -95,22 +92,18 @@ 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: plt.Axes, **plot_kwargs: Any) -> None:
         """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.
         """
-        fig = plt.figure()
-        ax = plt.gca()
         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')
 
 
 class TomographyResult:
diff --git a/examples/cross_entropy_benchmarking_example.py b/examples/cross_entropy_benchmarking_example.py
index a156a5f5888..d02821a766d 100644
--- a/examples/cross_entropy_benchmarking_example.py
+++ b/examples/cross_entropy_benchmarking_example.py
@@ -9,6 +9,8 @@
 details of this experiments.
 """
 
+import matplotlib as mpl
+
 import cirq
 
 
@@ -39,7 +41,9 @@ def main(repetitions=5000, num_circuits=20, cycles=range(2, 103, 10)):
         repetitions=repetitions)
 
     # Plot XEB fidelity vs number of cycles.
-    xeb_result.plot()
+    figure = mpl.figure.Figure()
+    ax = figure.add_subplot(111)
+    xeb_result.plot(ax)
 
 
 if __name__ == '__main__':
diff --git a/examples/qubit_characterizations_example.py b/examples/qubit_characterizations_example.py
index 0568207457d..a161f430c6b 100644
--- a/examples/qubit_characterizations_example.py
+++ b/examples/qubit_characterizations_example.py
@@ -1,3 +1,4 @@
+import matplotlib as mpl
 import numpy as np
 import cirq
 
@@ -10,21 +11,24 @@ def main():
     q_0 = cirq.GridQubit(0, 0)
     q_1 = cirq.GridQubit(0, 1)
 
+    figure = mpl.figure.Figure()
+    ax = figure.add_subplot(111)
+
     # Measure Rabi oscillation of q_0.
     rabi_results = cirq.experiments.rabi_oscillations(simulator, q_0, 4 * np.pi)
-    rabi_results.plot()
+    rabi_results.plot(ax)
 
     num_cfds = range(5, 20, 5)
 
     # Clifford-based randomized benchmarking of single-qubit gates on q_0.
     rb_result_1q = cirq.experiments.single_qubit_randomized_benchmarking(
         simulator, q_0, num_clifford_range=num_cfds, repetitions=100)
-    rb_result_1q.plot()
+    rb_result_1q.plot(ax)
 
     # Clifford-based randomized benchmarking of two-qubit gates on q_0 and q_1.
     rb_result_2q = cirq.experiments.two_qubit_randomized_benchmarking(
         simulator, q_0, q_1, num_clifford_range=num_cfds, repetitions=100)
-    rb_result_2q.plot()
+    rb_result_2q.plot(ax)
 
     # State-tomography of q_0 after application of an X/2 rotation.
     cir_1q = cirq.Circuit(cirq.X(q_0)**0.5)

From 958a01e82e280de59e2f2776e684f127e1885dc4 Mon Sep 17 00:00:00 2001
From: Ping Yeh <pingyeh@google.com>
Date: Wed, 9 Oct 2019 00:30:27 -0700
Subject: [PATCH 2/7] lint.

---
 cirq/experiments/cross_entropy_benchmarking_test.py | 2 +-
 examples/qubit_characterizations_example.py         | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/cirq/experiments/cross_entropy_benchmarking_test.py b/cirq/experiments/cross_entropy_benchmarking_test.py
index 46146db8710..fc92773d60d 100644
--- a/cirq/experiments/cross_entropy_benchmarking_test.py
+++ b/cirq/experiments/cross_entropy_benchmarking_test.py
@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import matplotlib as mpl
 import numpy as np
 
+import matplotlib as mpl
 import cirq
 
 from cirq.experiments import cross_entropy_benchmarking, build_entangling_layers
diff --git a/examples/qubit_characterizations_example.py b/examples/qubit_characterizations_example.py
index a161f430c6b..338e5a0e522 100644
--- a/examples/qubit_characterizations_example.py
+++ b/examples/qubit_characterizations_example.py
@@ -1,5 +1,6 @@
-import matplotlib as mpl
 import numpy as np
+
+import matplotlib as mpl
 import cirq
 
 

From 4d99fad848e5bb5b18b8dc276d9569229d8fbe62 Mon Sep 17 00:00:00 2001
From: Ping Yeh <pingyeh@google.com>
Date: Thu, 10 Oct 2019 15:48:39 -0700
Subject: [PATCH 3/7] Migrate plot methods to conform to docs/dev/plotting.md.

---
 .../experiments/cross_entropy_benchmarking.py |  18 ++-
 .../cross_entropy_benchmarking_test.py        |   5 +-
 cirq/experiments/qubit_characterizations.py   | 124 +++++++++++-------
 .../qubit_characterizations_test.py           |  15 +++
 dev_tools/incremental_coverage.py             |   2 +
 docs/dev/plotting.md                          |  11 +-
 .../cross_entropy_benchmarking_example.py     |   6 +-
 examples/qubit_characterizations_example.py   |  10 +-
 8 files changed, 120 insertions(+), 71 deletions(-)

diff --git a/cirq/experiments/cross_entropy_benchmarking.py b/cirq/experiments/cross_entropy_benchmarking.py
index 4d6031815de..3affed78802 100644
--- a/cirq/experiments/cross_entropy_benchmarking.py
+++ b/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,18 +44,28 @@ def data(self) -> Sequence[CrossEntropyPair]:
         """
         return self._data
 
-    def plot(self, ax: plt.Axes, **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]
         ax.set_ylim([0, 1.1])
         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(
diff --git a/cirq/experiments/cross_entropy_benchmarking_test.py b/cirq/experiments/cross_entropy_benchmarking_test.py
index fc92773d60d..7db102fd88c 100644
--- a/cirq/experiments/cross_entropy_benchmarking_test.py
+++ b/cirq/experiments/cross_entropy_benchmarking_test.py
@@ -14,7 +14,7 @@
 
 import numpy as np
 
-import matplotlib as mpl
+import matplotlib.pyplot as plt
 import cirq
 
 from cirq.experiments import cross_entropy_benchmarking, build_entangling_layers
@@ -79,6 +79,5 @@ def test_cross_entropy_benchmarking():
     assert len(fidelities_3) == 1
 
     # Sanity test that plot runs.
-    figure = mpl.figure.Figure()
-    ax = figure.add_subplot(111)
+    ax = plt.subplot()
     results_1.plot(ax)
diff --git a/cirq/experiments/qubit_characterizations.py b/cirq/experiments/qubit_characterizations.py
index cdfa0ca197e..aac4cf481a9 100644
--- a/cirq/experiments/qubit_characterizations.py
+++ b/cirq/experiments/qubit_characterizations.py
@@ -14,7 +14,7 @@
 
 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
 
@@ -54,18 +54,28 @@ 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, ax: plt.Axes, **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.
         """
+        show_plot = not ax
+        if not ax:
+            fig, ax = plt.subplots(1, 1, figsize=(8, 8))
         ax.set_ylim([0, 1])
         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:
@@ -92,18 +102,27 @@ 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, ax: plt.Axes, **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.
         """
+        show_plot = not ax
+        if not ax:
+            fig, ax = plt.subplots(1, 1, figsize=(8, 8))
         ax.set_ylim([0, 1])
-
         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:
@@ -123,12 +142,53 @@ 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,
@@ -482,11 +542,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()
@@ -495,49 +555,19 @@ 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)
+    ax.bar3d(x_indices, y_indices, z_indices, dx, dy, dz, color='#ff0080',
+             alpha=1.0, **bar3d_kwargs)
 
-    ax1.set_zlabel(z_label)
-    ax1.set_zlim3d(ylim[0], ylim[1])
+    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,
diff --git a/cirq/experiments/qubit_characterizations_test.py b/cirq/experiments/qubit_characterizations_test.py
index dec2888dbf9..aec5977f655 100644
--- a/cirq/experiments/qubit_characterizations_test.py
+++ b/cirq/experiments/qubit_characterizations_test.py
@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import matplotlib.pyplot as plt
 import numpy as np
+import pytest
 
 from cirq import GridQubit
 from cirq import circuits, ops, sim
@@ -142,3 +144,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)
diff --git a/dev_tools/incremental_coverage.py b/dev_tools/incremental_coverage.py
index 18628cd6867..10ef9dc0b8c 100644
--- a/dev_tools/incremental_coverage.py
+++ b/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'
 
diff --git a/docs/dev/plotting.md b/docs/dev/plotting.md
index 57257ebebaf..239fcfc036d 100644
--- a/docs/dev/plotting.md
+++ b/docs/dev/plotting.md
@@ -32,9 +32,9 @@ 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.
         ...
@@ -77,10 +77,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.
diff --git a/examples/cross_entropy_benchmarking_example.py b/examples/cross_entropy_benchmarking_example.py
index d02821a766d..a156a5f5888 100644
--- a/examples/cross_entropy_benchmarking_example.py
+++ b/examples/cross_entropy_benchmarking_example.py
@@ -9,8 +9,6 @@
 details of this experiments.
 """
 
-import matplotlib as mpl
-
 import cirq
 
 
@@ -41,9 +39,7 @@ def main(repetitions=5000, num_circuits=20, cycles=range(2, 103, 10)):
         repetitions=repetitions)
 
     # Plot XEB fidelity vs number of cycles.
-    figure = mpl.figure.Figure()
-    ax = figure.add_subplot(111)
-    xeb_result.plot(ax)
+    xeb_result.plot()
 
 
 if __name__ == '__main__':
diff --git a/examples/qubit_characterizations_example.py b/examples/qubit_characterizations_example.py
index 338e5a0e522..e561d2b01c5 100644
--- a/examples/qubit_characterizations_example.py
+++ b/examples/qubit_characterizations_example.py
@@ -1,6 +1,5 @@
 import numpy as np
 
-import matplotlib as mpl
 import cirq
 
 
@@ -12,24 +11,21 @@ def main():
     q_0 = cirq.GridQubit(0, 0)
     q_1 = cirq.GridQubit(0, 1)
 
-    figure = mpl.figure.Figure()
-    ax = figure.add_subplot(111)
-
     # Measure Rabi oscillation of q_0.
     rabi_results = cirq.experiments.rabi_oscillations(simulator, q_0, 4 * np.pi)
-    rabi_results.plot(ax)
+    rabi_results.plot()
 
     num_cfds = range(5, 20, 5)
 
     # Clifford-based randomized benchmarking of single-qubit gates on q_0.
     rb_result_1q = cirq.experiments.single_qubit_randomized_benchmarking(
         simulator, q_0, num_clifford_range=num_cfds, repetitions=100)
-    rb_result_1q.plot(ax)
+    rb_result_1q.plot()
 
     # Clifford-based randomized benchmarking of two-qubit gates on q_0 and q_1.
     rb_result_2q = cirq.experiments.two_qubit_randomized_benchmarking(
         simulator, q_0, q_1, num_clifford_range=num_cfds, repetitions=100)
-    rb_result_2q.plot(ax)
+    rb_result_2q.plot()
 
     # State-tomography of q_0 after application of an X/2 rotation.
     cir_1q = cirq.Circuit(cirq.X(q_0)**0.5)

From 0930a9f02e78526a39ccb19903e08da282c76c83 Mon Sep 17 00:00:00 2001
From: Ping Yeh <pingyeh@google.com>
Date: Thu, 10 Oct 2019 15:52:02 -0700
Subject: [PATCH 4/7] Drop unwanted changes in
 examples/qubit_characterizations_example.py.

---
 examples/qubit_characterizations_example.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/qubit_characterizations_example.py b/examples/qubit_characterizations_example.py
index e561d2b01c5..0568207457d 100644
--- a/examples/qubit_characterizations_example.py
+++ b/examples/qubit_characterizations_example.py
@@ -1,5 +1,4 @@
 import numpy as np
-
 import cirq
 
 

From 5ba53540164dc9ac43adbb60316db66c149aebb9 Mon Sep 17 00:00:00 2001
From: Ping Yeh <pingyeh@google.com>
Date: Thu, 10 Oct 2019 18:05:49 -0700
Subject: [PATCH 5/7] Format & lint.

---
 .../experiments/cross_entropy_benchmarking.py |  3 ++-
 cirq/experiments/qubit_characterizations.py   | 23 ++++++++++++++-----
 .../qubit_characterizations_test.py           |  3 ++-
 docs/dev/plotting.md                          |  2 +-
 4 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/cirq/experiments/cross_entropy_benchmarking.py b/cirq/experiments/cross_entropy_benchmarking.py
index 3affed78802..d0a96fd9db7 100644
--- a/cirq/experiments/cross_entropy_benchmarking.py
+++ b/cirq/experiments/cross_entropy_benchmarking.py
@@ -44,7 +44,8 @@ def data(self) -> Sequence[CrossEntropyPair]:
         """
         return self._data
 
-    def plot(self, ax: Optional[plt.Axes]=None, **plot_kwargs: Any) -> plt.Axes:
+    def plot(self, ax: Optional[plt.Axes] = None,
+             **plot_kwargs: Any) -> plt.Axes:
         """Plots the average XEB fidelity vs the number of cycles.
 
         Args:
diff --git a/cirq/experiments/qubit_characterizations.py b/cirq/experiments/qubit_characterizations.py
index aac4cf481a9..09186b19c90 100644
--- a/cirq/experiments/qubit_characterizations.py
+++ b/cirq/experiments/qubit_characterizations.py
@@ -54,7 +54,8 @@ 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, ax: Optional[plt.Axes]=None, **plot_kwargs: Any) -> plt.Axes:
+    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).
 
@@ -102,7 +103,8 @@ 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, ax: Optional[plt.Axes]=None, **plot_kwargs: Any) -> plt.Axes:
+    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.
 
@@ -142,7 +144,7 @@ def data(self) -> np.ndarray:
         """
         return self._density_matrix
 
-    def plot(self, axes: Optional[List[plt.Axes]]=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.
@@ -159,7 +161,9 @@ def plot(self, axes: Optional[List[plt.Axes]]=None,
         """
         show_plot = axes is None
         if axes is None:
-            fig, axes = plt.subplots(1, 2, figsize=(12.0, 5.0),
+            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.')
@@ -556,8 +560,15 @@ def _matrix_bar_plot(mat: np.ndarray,
     dx = np.ones(mat.size) * 0.3
     dy = np.ones(mat.size) * 0.3
     dz = mat.flatten()
-    ax.bar3d(x_indices, y_indices, z_indices, dx, dy, dz, color='#ff0080',
-             alpha=1.0, **bar3d_kwargs)
+    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])
diff --git a/cirq/experiments/qubit_characterizations_test.py b/cirq/experiments/qubit_characterizations_test.py
index aec5977f655..0ac83c8a19d 100644
--- a/cirq/experiments/qubit_characterizations_test.py
+++ b/cirq/experiments/qubit_characterizations_test.py
@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import matplotlib.pyplot as plt
 import numpy as np
 import pytest
 
+import matplotlib.pyplot as plt
+
 from cirq import GridQubit
 from cirq import circuits, ops, sim
 from cirq.experiments import (rabi_oscillations,
diff --git a/docs/dev/plotting.md b/docs/dev/plotting.md
index fcc5c3892f4..42cd6f8abd3 100644
--- a/docs/dev/plotting.md
+++ b/docs/dev/plotting.md
@@ -28,7 +28,7 @@ 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:

From 31a57663d63ac852b2d9a0acb921d89086853ad5 Mon Sep 17 00:00:00 2001
From: Ping Yeh <pingyeh@google.com>
Date: Thu, 10 Oct 2019 18:36:39 -0700
Subject: [PATCH 6/7] Remove a line that is seemingly dead code but fails
 pylint.

---
 cirq/experiments/qubit_characterizations.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cirq/experiments/qubit_characterizations.py b/cirq/experiments/qubit_characterizations.py
index 09186b19c90..e3395a5df0d 100644
--- a/cirq/experiments/qubit_characterizations.py
+++ b/cirq/experiments/qubit_characterizations.py
@@ -19,7 +19,6 @@
 import sympy
 
 from matplotlib import pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D  # type: ignore
 from cirq import circuits, devices, ops, protocols, study, work
 
 Cliffords = NamedTuple('Cliffords',

From 66e194c83a1acb6949c7a3bca331d6d61da24675 Mon Sep 17 00:00:00 2001
From: Ping Yeh <pingyeh@google.com>
Date: Thu, 10 Oct 2019 19:02:01 -0700
Subject: [PATCH 7/7] Add back the line with pylint annotation.

---
 cirq/experiments/qubit_characterizations.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cirq/experiments/qubit_characterizations.py b/cirq/experiments/qubit_characterizations.py
index e3395a5df0d..9eb0b077150 100644
--- a/cirq/experiments/qubit_characterizations.py
+++ b/cirq/experiments/qubit_characterizations.py
@@ -19,6 +19,7 @@
 import sympy
 
 from matplotlib import pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D  # type: ignore # pylint: disable=unused-import
 from cirq import circuits, devices, ops, protocols, study, work
 
 Cliffords = NamedTuple('Cliffords',