We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add to heatmap visualization tests
In the `test_colorbar` test there is a comment about testing that the position size and pad arguments are respected.




--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `cirq/vis/heatmap.py`
Content:
```
1 # Copyright 2019 The Cirq Developers
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """Heatmap class.
15 
16 See examples/bristlecone_heatmap_example.py for an example usage in
17 an interactive session.
18 """
19 
20 from typing import Any, Dict, List, Mapping, Optional, SupportsFloat, Tuple, Union
21 
22 import numpy as np
23 import matplotlib as mpl
24 import matplotlib.pyplot as plt
25 import pandas as pd
26 from matplotlib import collections as mpl_collections
27 from mpl_toolkits import axes_grid1
28 
29 from cirq.devices import grid_qubit
30 
31 QubitCoordinate = Union[Tuple[int, int], grid_qubit.GridQubit]
32 
33 # The value map is qubit coordinate -> a type that supports float conversion.
34 ValueMap = Union[Dict[grid_qubit.GridQubit, SupportsFloat], Dict[Tuple[int, int], SupportsFloat]]
35 
36 
37 def _get_qubit_row_col(qubit: QubitCoordinate) -> Tuple[int, int]:
38     if isinstance(qubit, grid_qubit.GridQubit):
39         return qubit.row, qubit.col
40     elif isinstance(qubit, tuple):
41         return qubit[0], qubit[1]
42 
43 
44 def relative_luminance(color: np.ndarray) -> float:
45     """Returns the relative luminance according to W3C specification.
46 
47     Spec: https://www.w3.org/TR/WCAG21/#dfn-relative-luminance.
48 
49     Args:
50         color: a numpy array with the first 3 elements red, green, and blue
51             with values in [0, 1].
52     Returns:
53         relative luminance of color in [0, 1].
54     """
55     rgb = color[:3]
56     rgb = np.where(rgb <= 0.03928, rgb / 12.92, ((rgb + 0.055) / 1.055) ** 2.4)
57     return rgb.dot([0.2126, 0.7152, 0.0722])
58 
59 
60 class Heatmap:
61     """Distribution of a value in 2D qubit lattice as a color map."""
62 
63     def __init__(self, value_map: ValueMap) -> None:
64         self.set_value_map(value_map)
65         self.annot_map = {  # Default annotation.
66             _get_qubit_row_col(qubit): format(float(value), '.2g')
67             for qubit, value in value_map.items()
68         }
69         self.annot_kwargs: Dict[str, Any] = {}
70         self.unset_url_map()
71         self.set_colorbar()
72         self.set_colormap()
73 
74     def set_value_map(self, value_map: ValueMap) -> 'Heatmap':
75         """Sets the values for each qubit.
76 
77         Args:
78             value_map: the values for determining color for each cell.
79         """
80         # Fail fast if float() fails.
81         # Keep the original value object for annotation.
82         self.value_map = {qubit: (float(value), value) for qubit, value in value_map.items()}
83         return self
84 
85     def set_annotation_map(
86         self, annot_map: Mapping[QubitCoordinate, str], **text_options: str
87     ) -> 'Heatmap':
88         """Sets the annotation text for each qubit.
89 
90         Note that set_annotation_map() and set_annotation_format()
91         both sets the annotation map to be used. Whichever is called later wins.
92 
93         Args:
94             annot_map: the texts to be drawn on each qubit cell.
95             text_options: keyword arguments passed to matplotlib.text.Text()
96                 when drawing the annotation texts.
97         """
98         self.annot_map = {_get_qubit_row_col(qubit): value for qubit, value in annot_map.items()}
99         self.annot_kwargs = text_options
100         return self
101 
102     def set_annotation_format(self, annot_format: str, **text_options: str) -> 'Heatmap':
103         """Sets a format string to format values for each qubit.
104 
105         Args:
106             annot_format: the format string for formatting values.
107             text_options: keyword arguments to matplotlib.text.Text().
108         """
109         self.annot_map = {
110             _get_qubit_row_col(qubit): format(value[1], annot_format)
111             for qubit, value in self.value_map.items()
112         }
113         self.annot_kwargs = text_options
114         return self
115 
116     def unset_annotation(self) -> 'Heatmap':
117         """Disables annotation. No texts are shown in cells."""
118         self.annot_map = {}
119         return self
120 
121     def set_url_map(self, url_map: Mapping[QubitCoordinate, str]) -> 'Heatmap':
122         """Sets the URLs for each cell."""
123         self.url_map = {_get_qubit_row_col(qubit): value for qubit, value in url_map.items()}
124         return self
125 
126     def unset_url_map(self) -> 'Heatmap':
127         """Disables URL. No URLs are associated with cells."""
128         self.url_map = {}
129         return self
130 
131     def set_colorbar(
132         self, position: str = 'right', size: str = '5%', pad: str = '2%', **colorbar_options: Any
133     ) -> 'Heatmap':
134         """Sets location and style of colorbar.
135 
136         Args:
137             position: colorbar position, one of 'left'|'right'|'top'|'bottom'.
138             size: a string ending in '%' to specify the width of the colorbar.
139                 Nominally, '100%' means the same width as the heatmap.
140             pad: a string ending in '%' to specify the space between the
141                 colorbar and the heatmap.
142             colorbar_options: keyword arguments passed to
143                 matplotlib.Figure.colorbar().
144         """
145         self.plot_colorbar = True
146         self.colorbar_location_options = {'position': position, 'size': size, 'pad': pad}
147         self.colorbar_options = colorbar_options
148         return self
149 
150     def unset_colorbar(self) -> 'Heatmap':
151         """Disables colorbar. No colorbar is drawn."""
152         self.plot_colorbar = False
153         return self
154 
155     def set_colormap(
156         self,
157         colormap: Union[str, mpl.colors.Colormap] = 'viridis',
158         vmin: Optional[float] = None,
159         vmax: Optional[float] = None,
160     ) -> 'Heatmap':
161         """Sets the colormap.
162 
163         Args:
164             colormap: either a colormap name or a Colormap instance.
165             vmin: the minimum value to map to the minimum color. Default is
166                 the minimum value in value_map.
167             vmax: the maximum value to map to the maximum color. Default is
168                 the maximum value in value_map.
169         """
170         self.colormap = colormap
171         self.vmin = vmin
172         self.vmax = vmax
173         return self
174 
175     def plot(
176         self, ax: Optional[plt.Axes] = None, **pcolor_options: Any
177     ) -> Tuple[plt.Axes, mpl_collections.Collection, pd.DataFrame]:
178         """Plots the heatmap on the given Axes.
179 
180         Args:
181             ax: the Axes to plot on. If not given, a new figure is created,
182                 plotted on, and shown.
183             pcolor_options: keyword arguments passed to ax.pcolor().
184 
185         Returns:
186             A 3-tuple ``(ax, mesh, value_table)``. ``ax`` is the `plt.Axes` that
187             is plotted on. ``mesh`` is the collection of paths drawn and filled.
188             ``value_table`` is the 2-D pandas DataFrame of values constructed
189             from the value_map.
190         """
191         show_plot = not ax
192         if not ax:
193             fig, ax = plt.subplots(figsize=(8, 8))
194         # Find the boundary and size of the heatmap.
195         coordinate_list = [_get_qubit_row_col(qubit) for qubit in self.value_map.keys()]
196         rows = [row for row, _ in coordinate_list]
197         cols = [col for _, col in coordinate_list]
198         min_row, max_row = min(rows), max(rows)
199         min_col, max_col = min(cols), max(cols)
200         height, width = max_row - min_row + 1, max_col - min_col + 1
201         # Construct the (height x width) table of values. Cells with no values
202         # are filled with np.nan.
203         value_table = pd.DataFrame(
204             np.nan, index=range(min_row, max_row + 1), columns=range(min_col, max_col + 1)
205         )
206         for qubit, (float_value, _) in self.value_map.items():
207             row, col = _get_qubit_row_col(qubit)
208             value_table[col][row] = float_value
209         # Construct the (height + 1) x (width + 1) cell boundary tables.
210         x_table = np.array([np.arange(min_col - 0.5, max_col + 1.5)] * (height + 1))
211         y_table = np.array([np.arange(min_row - 0.5, max_row + 1.5)] * (width + 1)).transpose()
212 
213         # Construct the URL array as an ordered list of URLs for non-nan cells.
214         url_array: List[str] = []
215         if self.url_map:
216             url_array = [self.url_map.get((row, col), '') for row, col in value_table.stack().index]
217 
218         # Plot the heatmap.
219         mesh = ax.pcolor(
220             x_table,
221             y_table,
222             value_table,
223             vmin=self.vmin,
224             vmax=self.vmax,
225             cmap=self.colormap,
226             urls=url_array,
227             **pcolor_options,
228         )
229         mesh.update_scalarmappable()
230         ax.set(xlabel='column', ylabel='row')
231         ax.xaxis.set_ticks(np.arange(min_col, max_col + 1))
232         ax.yaxis.set_ticks(np.arange(min_row, max_row + 1))
233         ax.set_ylim((max_row + 0.5, min_row - 0.5))
234 
235         if self.plot_colorbar:
236             self._plot_colorbar(mesh, ax)
237 
238         if self.annot_map:
239             self._write_annotations(mesh, ax)
240 
241         if show_plot:
242             fig.show()
243 
244         return ax, mesh, value_table
245 
246     def _plot_colorbar(
247         self, mappable: mpl.cm.ScalarMappable, ax: plt.Axes
248     ) -> mpl.colorbar.Colorbar:
249         """Plots the colorbar. Internal."""
250         colorbar_ax = axes_grid1.make_axes_locatable(ax).append_axes(
251             **self.colorbar_location_options
252         )
253         position = self.colorbar_location_options.get('position', 'right')
254         orien = 'vertical' if position in ('left', 'right') else 'horizontal'
255         colorbar = ax.figure.colorbar(
256             mappable, colorbar_ax, ax, orientation=orien, **self.colorbar_options
257         )
258         colorbar_ax.tick_params(axis='y', direction='out')
259         return colorbar
260 
261     def _write_annotations(self, mesh: mpl_collections.Collection, ax: plt.Axes) -> None:
262         """Writes annotations to the center of cells. Internal."""
263         for path, facecolor in zip(mesh.get_paths(), mesh.get_facecolors()):
264             # Calculate the center of the cell, assuming that it is a square
265             # centered at (x=col, y=row).
266             vertices = path.vertices[:4]
267             row = int(round(np.mean([v[1] for v in vertices])))
268             col = int(round(np.mean([v[0] for v in vertices])))
269             annotation = self.annot_map.get((row, col), '')
270             if not annotation:
271                 continue
272             face_luminance = relative_luminance(facecolor)
273             text_color = 'black' if face_luminance > 0.4 else 'white'
274             text_kwargs = dict(color=text_color, ha="center", va="center")
275             text_kwargs.update(self.annot_kwargs)
276             ax.text(col, row, annotation, **text_kwargs)
277
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/cirq/vis/heatmap.py b/cirq/vis/heatmap.py
--- a/cirq/vis/heatmap.py
+++ b/cirq/vis/heatmap.py
@@ -38,7 +38,7 @@
     if isinstance(qubit, grid_qubit.GridQubit):
         return qubit.row, qubit.col
     elif isinstance(qubit, tuple):
-        return qubit[0], qubit[1]
+        return int(qubit[0]), int(qubit[1])
 
 
 def relative_luminance(color: np.ndarray) -> float:

{"golden_diff": "diff --git a/cirq/vis/heatmap.py b/cirq/vis/heatmap.py\n--- a/cirq/vis/heatmap.py\n+++ b/cirq/vis/heatmap.py\n@@ -38,7 +38,7 @@\n     if isinstance(qubit, grid_qubit.GridQubit):\n         return qubit.row, qubit.col\n     elif isinstance(qubit, tuple):\n-        return qubit[0], qubit[1]\n+        return int(qubit[0]), int(qubit[1])\n \n \n def relative_luminance(color: np.ndarray) -> float:\n", "issue": "Add to heatmap visualization tests\nIn the `test_colorbar` test there is a comment about testing that the position size and pad arguments are respected.\r\n\r\n\r\n\n", "before_files": [{"content": "# Copyright 2019 The Cirq Developers\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Heatmap class.\n\nSee examples/bristlecone_heatmap_example.py for an example usage in\nan interactive session.\n\"\"\"\n\nfrom typing import Any, Dict, List, Mapping, Optional, SupportsFloat, Tuple, Union\n\nimport numpy as np\nimport matplotlib as mpl\nimport matplotlib.pyplot as plt\nimport pandas as pd\nfrom matplotlib import collections as mpl_collections\nfrom mpl_toolkits import axes_grid1\n\nfrom cirq.devices import grid_qubit\n\nQubitCoordinate = Union[Tuple[int, int], grid_qubit.GridQubit]\n\n# The value map is qubit coordinate -> a type that supports float conversion.\nValueMap = Union[Dict[grid_qubit.GridQubit, SupportsFloat], Dict[Tuple[int, int], SupportsFloat]]\n\n\ndef _get_qubit_row_col(qubit: QubitCoordinate) -> Tuple[int, int]:\n    if isinstance(qubit, grid_qubit.GridQubit):\n        return qubit.row, qubit.col\n    elif isinstance(qubit, tuple):\n        return qubit[0], qubit[1]\n\n\ndef relative_luminance(color: np.ndarray) -> float:\n    \"\"\"Returns the relative luminance according to W3C specification.\n\n    Spec: https://www.w3.org/TR/WCAG21/#dfn-relative-luminance.\n\n    Args:\n        color: a numpy array with the first 3 elements red, green, and blue\n            with values in [0, 1].\n    Returns:\n        relative luminance of color in [0, 1].\n    \"\"\"\n    rgb = color[:3]\n    rgb = np.where(rgb <= 0.03928, rgb / 12.92, ((rgb + 0.055) / 1.055) ** 2.4)\n    return rgb.dot([0.2126, 0.7152, 0.0722])\n\n\nclass Heatmap:\n    \"\"\"Distribution of a value in 2D qubit lattice as a color map.\"\"\"\n\n    def __init__(self, value_map: ValueMap) -> None:\n        self.set_value_map(value_map)\n        self.annot_map = {  # Default annotation.\n            _get_qubit_row_col(qubit): format(float(value), '.2g')\n            for qubit, value in value_map.items()\n        }\n        self.annot_kwargs: Dict[str, Any] = {}\n        self.unset_url_map()\n        self.set_colorbar()\n        self.set_colormap()\n\n    def set_value_map(self, value_map: ValueMap) -> 'Heatmap':\n        \"\"\"Sets the values for each qubit.\n\n        Args:\n            value_map: the values for determining color for each cell.\n        \"\"\"\n        # Fail fast if float() fails.\n        # Keep the original value object for annotation.\n        self.value_map = {qubit: (float(value), value) for qubit, value in value_map.items()}\n        return self\n\n    def set_annotation_map(\n        self, annot_map: Mapping[QubitCoordinate, str], **text_options: str\n    ) -> 'Heatmap':\n        \"\"\"Sets the annotation text for each qubit.\n\n        Note that set_annotation_map() and set_annotation_format()\n        both sets the annotation map to be used. Whichever is called later wins.\n\n        Args:\n            annot_map: the texts to be drawn on each qubit cell.\n            text_options: keyword arguments passed to matplotlib.text.Text()\n                when drawing the annotation texts.\n        \"\"\"\n        self.annot_map = {_get_qubit_row_col(qubit): value for qubit, value in annot_map.items()}\n        self.annot_kwargs = text_options\n        return self\n\n    def set_annotation_format(self, annot_format: str, **text_options: str) -> 'Heatmap':\n        \"\"\"Sets a format string to format values for each qubit.\n\n        Args:\n            annot_format: the format string for formatting values.\n            text_options: keyword arguments to matplotlib.text.Text().\n        \"\"\"\n        self.annot_map = {\n            _get_qubit_row_col(qubit): format(value[1], annot_format)\n            for qubit, value in self.value_map.items()\n        }\n        self.annot_kwargs = text_options\n        return self\n\n    def unset_annotation(self) -> 'Heatmap':\n        \"\"\"Disables annotation. No texts are shown in cells.\"\"\"\n        self.annot_map = {}\n        return self\n\n    def set_url_map(self, url_map: Mapping[QubitCoordinate, str]) -> 'Heatmap':\n        \"\"\"Sets the URLs for each cell.\"\"\"\n        self.url_map = {_get_qubit_row_col(qubit): value for qubit, value in url_map.items()}\n        return self\n\n    def unset_url_map(self) -> 'Heatmap':\n        \"\"\"Disables URL. No URLs are associated with cells.\"\"\"\n        self.url_map = {}\n        return self\n\n    def set_colorbar(\n        self, position: str = 'right', size: str = '5%', pad: str = '2%', **colorbar_options: Any\n    ) -> 'Heatmap':\n        \"\"\"Sets location and style of colorbar.\n\n        Args:\n            position: colorbar position, one of 'left'|'right'|'top'|'bottom'.\n            size: a string ending in '%' to specify the width of the colorbar.\n                Nominally, '100%' means the same width as the heatmap.\n            pad: a string ending in '%' to specify the space between the\n                colorbar and the heatmap.\n            colorbar_options: keyword arguments passed to\n                matplotlib.Figure.colorbar().\n        \"\"\"\n        self.plot_colorbar = True\n        self.colorbar_location_options = {'position': position, 'size': size, 'pad': pad}\n        self.colorbar_options = colorbar_options\n        return self\n\n    def unset_colorbar(self) -> 'Heatmap':\n        \"\"\"Disables colorbar. No colorbar is drawn.\"\"\"\n        self.plot_colorbar = False\n        return self\n\n    def set_colormap(\n        self,\n        colormap: Union[str, mpl.colors.Colormap] = 'viridis',\n        vmin: Optional[float] = None,\n        vmax: Optional[float] = None,\n    ) -> 'Heatmap':\n        \"\"\"Sets the colormap.\n\n        Args:\n            colormap: either a colormap name or a Colormap instance.\n            vmin: the minimum value to map to the minimum color. Default is\n                the minimum value in value_map.\n            vmax: the maximum value to map to the maximum color. Default is\n                the maximum value in value_map.\n        \"\"\"\n        self.colormap = colormap\n        self.vmin = vmin\n        self.vmax = vmax\n        return self\n\n    def plot(\n        self, ax: Optional[plt.Axes] = None, **pcolor_options: Any\n    ) -> Tuple[plt.Axes, mpl_collections.Collection, pd.DataFrame]:\n        \"\"\"Plots the heatmap on the given Axes.\n\n        Args:\n            ax: the Axes to plot on. If not given, a new figure is created,\n                plotted on, and shown.\n            pcolor_options: keyword arguments passed to ax.pcolor().\n\n        Returns:\n            A 3-tuple ``(ax, mesh, value_table)``. ``ax`` is the `plt.Axes` that\n            is plotted on. ``mesh`` is the collection of paths drawn and filled.\n            ``value_table`` is the 2-D pandas DataFrame of values constructed\n            from the value_map.\n        \"\"\"\n        show_plot = not ax\n        if not ax:\n            fig, ax = plt.subplots(figsize=(8, 8))\n        # Find the boundary and size of the heatmap.\n        coordinate_list = [_get_qubit_row_col(qubit) for qubit in self.value_map.keys()]\n        rows = [row for row, _ in coordinate_list]\n        cols = [col for _, col in coordinate_list]\n        min_row, max_row = min(rows), max(rows)\n        min_col, max_col = min(cols), max(cols)\n        height, width = max_row - min_row + 1, max_col - min_col + 1\n        # Construct the (height x width) table of values. Cells with no values\n        # are filled with np.nan.\n        value_table = pd.DataFrame(\n            np.nan, index=range(min_row, max_row + 1), columns=range(min_col, max_col + 1)\n        )\n        for qubit, (float_value, _) in self.value_map.items():\n            row, col = _get_qubit_row_col(qubit)\n            value_table[col][row] = float_value\n        # Construct the (height + 1) x (width + 1) cell boundary tables.\n        x_table = np.array([np.arange(min_col - 0.5, max_col + 1.5)] * (height + 1))\n        y_table = np.array([np.arange(min_row - 0.5, max_row + 1.5)] * (width + 1)).transpose()\n\n        # Construct the URL array as an ordered list of URLs for non-nan cells.\n        url_array: List[str] = []\n        if self.url_map:\n            url_array = [self.url_map.get((row, col), '') for row, col in value_table.stack().index]\n\n        # Plot the heatmap.\n        mesh = ax.pcolor(\n            x_table,\n            y_table,\n            value_table,\n            vmin=self.vmin,\n            vmax=self.vmax,\n            cmap=self.colormap,\n            urls=url_array,\n            **pcolor_options,\n        )\n        mesh.update_scalarmappable()\n        ax.set(xlabel='column', ylabel='row')\n        ax.xaxis.set_ticks(np.arange(min_col, max_col + 1))\n        ax.yaxis.set_ticks(np.arange(min_row, max_row + 1))\n        ax.set_ylim((max_row + 0.5, min_row - 0.5))\n\n        if self.plot_colorbar:\n            self._plot_colorbar(mesh, ax)\n\n        if self.annot_map:\n            self._write_annotations(mesh, ax)\n\n        if show_plot:\n            fig.show()\n\n        return ax, mesh, value_table\n\n    def _plot_colorbar(\n        self, mappable: mpl.cm.ScalarMappable, ax: plt.Axes\n    ) -> mpl.colorbar.Colorbar:\n        \"\"\"Plots the colorbar. Internal.\"\"\"\n        colorbar_ax = axes_grid1.make_axes_locatable(ax).append_axes(\n            **self.colorbar_location_options\n        )\n        position = self.colorbar_location_options.get('position', 'right')\n        orien = 'vertical' if position in ('left', 'right') else 'horizontal'\n        colorbar = ax.figure.colorbar(\n            mappable, colorbar_ax, ax, orientation=orien, **self.colorbar_options\n        )\n        colorbar_ax.tick_params(axis='y', direction='out')\n        return colorbar\n\n    def _write_annotations(self, mesh: mpl_collections.Collection, ax: plt.Axes) -> None:\n        \"\"\"Writes annotations to the center of cells. Internal.\"\"\"\n        for path, facecolor in zip(mesh.get_paths(), mesh.get_facecolors()):\n            # Calculate the center of the cell, assuming that it is a square\n            # centered at (x=col, y=row).\n            vertices = path.vertices[:4]\n            row = int(round(np.mean([v[1] for v in vertices])))\n            col = int(round(np.mean([v[0] for v in vertices])))\n            annotation = self.annot_map.get((row, col), '')\n            if not annotation:\n                continue\n            face_luminance = relative_luminance(facecolor)\n            text_color = 'black' if face_luminance > 0.4 else 'white'\n            text_kwargs = dict(color=text_color, ha=\"center\", va=\"center\")\n            text_kwargs.update(self.annot_kwargs)\n            ax.text(col, row, annotation, **text_kwargs)\n", "path": "cirq/vis/heatmap.py"}], "after_files": [{"content": "# Copyright 2019 The Cirq Developers\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Heatmap class.\n\nSee examples/bristlecone_heatmap_example.py for an example usage in\nan interactive session.\n\"\"\"\n\nfrom typing import Any, Dict, List, Mapping, Optional, SupportsFloat, Tuple, Union\n\nimport numpy as np\nimport matplotlib as mpl\nimport matplotlib.pyplot as plt\nimport pandas as pd\nfrom matplotlib import collections as mpl_collections\nfrom mpl_toolkits import axes_grid1\n\nfrom cirq.devices import grid_qubit\n\nQubitCoordinate = Union[Tuple[int, int], grid_qubit.GridQubit]\n\n# The value map is qubit coordinate -> a type that supports float conversion.\nValueMap = Union[Dict[grid_qubit.GridQubit, SupportsFloat], Dict[Tuple[int, int], SupportsFloat]]\n\n\ndef _get_qubit_row_col(qubit: QubitCoordinate) -> Tuple[int, int]:\n    if isinstance(qubit, grid_qubit.GridQubit):\n        return qubit.row, qubit.col\n    elif isinstance(qubit, tuple):\n        return int(qubit[0]), int(qubit[1])\n\n\ndef relative_luminance(color: np.ndarray) -> float:\n    \"\"\"Returns the relative luminance according to W3C specification.\n\n    Spec: https://www.w3.org/TR/WCAG21/#dfn-relative-luminance.\n\n    Args:\n        color: a numpy array with the first 3 elements red, green, and blue\n            with values in [0, 1].\n    Returns:\n        relative luminance of color in [0, 1].\n    \"\"\"\n    rgb = color[:3]\n    rgb = np.where(rgb <= 0.03928, rgb / 12.92, ((rgb + 0.055) / 1.055) ** 2.4)\n    return rgb.dot([0.2126, 0.7152, 0.0722])\n\n\nclass Heatmap:\n    \"\"\"Distribution of a value in 2D qubit lattice as a color map.\"\"\"\n\n    def __init__(self, value_map: ValueMap) -> None:\n        self.set_value_map(value_map)\n        self.annot_map = {  # Default annotation.\n            _get_qubit_row_col(qubit): format(float(value), '.2g')\n            for qubit, value in value_map.items()\n        }\n        self.annot_kwargs: Dict[str, Any] = {}\n        self.unset_url_map()\n        self.set_colorbar()\n        self.set_colormap()\n\n    def set_value_map(self, value_map: ValueMap) -> 'Heatmap':\n        \"\"\"Sets the values for each qubit.\n\n        Args:\n            value_map: the values for determining color for each cell.\n        \"\"\"\n        # Fail fast if float() fails.\n        # Keep the original value object for annotation.\n        self.value_map = {qubit: (float(value), value) for qubit, value in value_map.items()}\n        return self\n\n    def set_annotation_map(\n        self, annot_map: Mapping[QubitCoordinate, str], **text_options: str\n    ) -> 'Heatmap':\n        \"\"\"Sets the annotation text for each qubit.\n\n        Note that set_annotation_map() and set_annotation_format()\n        both sets the annotation map to be used. Whichever is called later wins.\n\n        Args:\n            annot_map: the texts to be drawn on each qubit cell.\n            text_options: keyword arguments passed to matplotlib.text.Text()\n                when drawing the annotation texts.\n        \"\"\"\n        self.annot_map = {_get_qubit_row_col(qubit): value for qubit, value in annot_map.items()}\n        self.annot_kwargs = text_options\n        return self\n\n    def set_annotation_format(self, annot_format: str, **text_options: str) -> 'Heatmap':\n        \"\"\"Sets a format string to format values for each qubit.\n\n        Args:\n            annot_format: the format string for formatting values.\n            text_options: keyword arguments to matplotlib.text.Text().\n        \"\"\"\n        self.annot_map = {\n            _get_qubit_row_col(qubit): format(value[1], annot_format)\n            for qubit, value in self.value_map.items()\n        }\n        self.annot_kwargs = text_options\n        return self\n\n    def unset_annotation(self) -> 'Heatmap':\n        \"\"\"Disables annotation. No texts are shown in cells.\"\"\"\n        self.annot_map = {}\n        return self\n\n    def set_url_map(self, url_map: Mapping[QubitCoordinate, str]) -> 'Heatmap':\n        \"\"\"Sets the URLs for each cell.\"\"\"\n        self.url_map = {_get_qubit_row_col(qubit): value for qubit, value in url_map.items()}\n        return self\n\n    def unset_url_map(self) -> 'Heatmap':\n        \"\"\"Disables URL. No URLs are associated with cells.\"\"\"\n        self.url_map = {}\n        return self\n\n    def set_colorbar(\n        self, position: str = 'right', size: str = '5%', pad: str = '2%', **colorbar_options: Any\n    ) -> 'Heatmap':\n        \"\"\"Sets location and style of colorbar.\n\n        Args:\n            position: colorbar position, one of 'left'|'right'|'top'|'bottom'.\n            size: a string ending in '%' to specify the width of the colorbar.\n                Nominally, '100%' means the same width as the heatmap.\n            pad: a string ending in '%' to specify the space between the\n                colorbar and the heatmap.\n            colorbar_options: keyword arguments passed to\n                matplotlib.Figure.colorbar().\n        \"\"\"\n        self.plot_colorbar = True\n        self.colorbar_location_options = {'position': position, 'size': size, 'pad': pad}\n        self.colorbar_options = colorbar_options\n        return self\n\n    def unset_colorbar(self) -> 'Heatmap':\n        \"\"\"Disables colorbar. No colorbar is drawn.\"\"\"\n        self.plot_colorbar = False\n        return self\n\n    def set_colormap(\n        self,\n        colormap: Union[str, mpl.colors.Colormap] = 'viridis',\n        vmin: Optional[float] = None,\n        vmax: Optional[float] = None,\n    ) -> 'Heatmap':\n        \"\"\"Sets the colormap.\n\n        Args:\n            colormap: either a colormap name or a Colormap instance.\n            vmin: the minimum value to map to the minimum color. Default is\n                the minimum value in value_map.\n            vmax: the maximum value to map to the maximum color. Default is\n                the maximum value in value_map.\n        \"\"\"\n        self.colormap = colormap\n        self.vmin = vmin\n        self.vmax = vmax\n        return self\n\n    def plot(\n        self, ax: Optional[plt.Axes] = None, **pcolor_options: Any\n    ) -> Tuple[plt.Axes, mpl_collections.Collection, pd.DataFrame]:\n        \"\"\"Plots the heatmap on the given Axes.\n\n        Args:\n            ax: the Axes to plot on. If not given, a new figure is created,\n                plotted on, and shown.\n            pcolor_options: keyword arguments passed to ax.pcolor().\n\n        Returns:\n            A 3-tuple ``(ax, mesh, value_table)``. ``ax`` is the `plt.Axes` that\n            is plotted on. ``mesh`` is the collection of paths drawn and filled.\n            ``value_table`` is the 2-D pandas DataFrame of values constructed\n            from the value_map.\n        \"\"\"\n        show_plot = not ax\n        if not ax:\n            fig, ax = plt.subplots(figsize=(8, 8))\n        # Find the boundary and size of the heatmap.\n        coordinate_list = [_get_qubit_row_col(qubit) for qubit in self.value_map.keys()]\n        rows = [row for row, _ in coordinate_list]\n        cols = [col for _, col in coordinate_list]\n        min_row, max_row = min(rows), max(rows)\n        min_col, max_col = min(cols), max(cols)\n        height, width = max_row - min_row + 1, max_col - min_col + 1\n        # Construct the (height x width) table of values. Cells with no values\n        # are filled with np.nan.\n        value_table = pd.DataFrame(\n            np.nan, index=range(min_row, max_row + 1), columns=range(min_col, max_col + 1)\n        )\n        for qubit, (float_value, _) in self.value_map.items():\n            row, col = _get_qubit_row_col(qubit)\n            value_table[col][row] = float_value\n        # Construct the (height + 1) x (width + 1) cell boundary tables.\n        x_table = np.array([np.arange(min_col - 0.5, max_col + 1.5)] * (height + 1))\n        y_table = np.array([np.arange(min_row - 0.5, max_row + 1.5)] * (width + 1)).transpose()\n\n        # Construct the URL array as an ordered list of URLs for non-nan cells.\n        url_array: List[str] = []\n        if self.url_map:\n            url_array = [self.url_map.get((row, col), '') for row, col in value_table.stack().index]\n\n        # Plot the heatmap.\n        mesh = ax.pcolor(\n            x_table,\n            y_table,\n            value_table,\n            vmin=self.vmin,\n            vmax=self.vmax,\n            cmap=self.colormap,\n            urls=url_array,\n            **pcolor_options,\n        )\n        mesh.update_scalarmappable()\n        ax.set(xlabel='column', ylabel='row')\n        ax.xaxis.set_ticks(np.arange(min_col, max_col + 1))\n        ax.yaxis.set_ticks(np.arange(min_row, max_row + 1))\n        ax.set_ylim((max_row + 0.5, min_row - 0.5))\n\n        if self.plot_colorbar:\n            self._plot_colorbar(mesh, ax)\n\n        if self.annot_map:\n            self._write_annotations(mesh, ax)\n\n        if show_plot:\n            fig.show()\n\n        return ax, mesh, value_table\n\n    def _plot_colorbar(\n        self, mappable: mpl.cm.ScalarMappable, ax: plt.Axes\n    ) -> mpl.colorbar.Colorbar:\n        \"\"\"Plots the colorbar. Internal.\"\"\"\n        colorbar_ax = axes_grid1.make_axes_locatable(ax).append_axes(\n            **self.colorbar_location_options\n        )\n        position = self.colorbar_location_options.get('position', 'right')\n        orien = 'vertical' if position in ('left', 'right') else 'horizontal'\n        colorbar = ax.figure.colorbar(\n            mappable, colorbar_ax, ax, orientation=orien, **self.colorbar_options\n        )\n        colorbar_ax.tick_params(axis='y', direction='out')\n        return colorbar\n\n    def _write_annotations(self, mesh: mpl_collections.Collection, ax: plt.Axes) -> None:\n        \"\"\"Writes annotations to the center of cells. Internal.\"\"\"\n        for path, facecolor in zip(mesh.get_paths(), mesh.get_facecolors()):\n            # Calculate the center of the cell, assuming that it is a square\n            # centered at (x=col, y=row).\n            vertices = path.vertices[:4]\n            row = int(round(np.mean([v[1] for v in vertices])))\n            col = int(round(np.mean([v[0] for v in vertices])))\n            annotation = self.annot_map.get((row, col), '')\n            if not annotation:\n                continue\n            face_luminance = relative_luminance(facecolor)\n            text_color = 'black' if face_luminance > 0.4 else 'white'\n            text_kwargs = dict(color=text_color, ha=\"center\", va=\"center\")\n            text_kwargs.update(self.annot_kwargs)\n            ax.text(col, row, annotation, **text_kwargs)\n", "path": "cirq/vis/heatmap.py"}]}