Repository: quantumlib/Cirq

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.





Discussion:
I think I can contribute to this issue
Looks like that more works besides the test are required. For instance:
```python
value_map = {
    (qubit.row, qubit.col): np.random.random() for qubit in cirq.google.Bristlecone.qubits
}

heatmap = cirq.Heatmap(value_map)
heatmap.set_colorbar(position='left') # specific the colorbar to left instead of right (allowed by current design)
heatmap.plot()
```

Will generate the overlap between colorbar and heatmap:
![image](https://user-images.githubusercontent.com/16711681/105103915-3064ab80-5a66-11eb-9e64-a21a3128b3d1.png)

Resolving patch:
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 @@ def _get_qubit_row_col(qubit: QubitCoordinate) -> Tuple[int, int]:
     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:


Test patch:
diff --git a/cirq/vis/heatmap_test.py b/cirq/vis/heatmap_test.py
--- a/cirq/vis/heatmap_test.py
+++ b/cirq/vis/heatmap_test.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 """Tests for Heatmap."""
 
+import pathlib
 import string
+from tempfile import mkdtemp
 
 import numpy as np
 import pytest
@@ -226,20 +228,72 @@ def test_urls(ax, test_GridQubit):
     assert mesh.get_urls() == expected_urls
 
 
-def test_colorbar(ax):
+@pytest.mark.parametrize(
+    'position,size,pad',
+    [
+        ('right', "5%", "2%"),
+        ('right', "5%", "10%"),
+        ('right', "20%", "2%"),
+        ('right', "20%", "10%"),
+        ('left', "5%", "2%"),
+        ('left', "5%", "10%"),
+        ('left', "20%", "2%"),
+        ('left', "20%", "10%"),
+        ('top', "5%", "2%"),
+        ('top', "5%", "10%"),
+        ('top', "20%", "2%"),
+        ('top', "20%", "10%"),
+        ('bottom', "5%", "2%"),
+        ('bottom', "5%", "10%"),
+        ('bottom', "20%", "2%"),
+        ('bottom', "20%", "10%"),
+    ],
+)
+def test_colorbar(ax, position, size, pad):
     qubits = ((0, 5), (8, 1), (7, 0), (13, 5), (1, 6), (3, 2), (2, 8))
     values = np.random.random(len(qubits))
     test_value_map = {qubit: value for qubit, value in zip(qubits, values)}
     random_heatmap = heatmap.Heatmap(test_value_map).unset_colorbar()
     fig1, ax1 = plt.subplots()
     random_heatmap.plot(ax1)
-    random_heatmap.set_colorbar()
+    random_heatmap.set_colorbar(position=position, size=size, pad=pad)
     fig2, ax2 = plt.subplots()
     random_heatmap.plot(ax2)
 
+    # We need to call savefig() explicitly for updating axes position since the figure
+    # object has been altered in the HeatMap._plot_colorbar function.
+    tmp_dir = mkdtemp()
+    fig2.savefig(pathlib.Path(tmp_dir) / 'tmp.png')
+
     # Check that the figure has one more object in it when colorbar is on.
     assert len(fig2.get_children()) == len(fig1.get_children()) + 1
 
-    # TODO: Make this is a more thorough test, e.g., we should test that the
-    # position, size and pad arguments are respected.
-    # Github issue: https://github.com/quantumlib/Cirq/issues/2969
+    fig_pos = fig2.get_axes()[0].get_position()
+    colorbar_pos = fig2.get_axes()[1].get_position()
+
+    origin_axes_size = (
+        fig_pos.xmax - fig_pos.xmin
+        if position in ["left", "right"]
+        else fig_pos.ymax - fig_pos.ymin
+    )
+    expected_pad = int(pad.replace("%", "")) / 100 * origin_axes_size
+    expected_size = int(size.replace("%", "")) / 100 * origin_axes_size
+
+    if position == "right":
+        pad_distance = colorbar_pos.xmin - fig_pos.xmax
+        colorbar_size = colorbar_pos.xmax - colorbar_pos.xmin
+    elif position == "left":
+        pad_distance = fig_pos.xmin - colorbar_pos.xmax
+        colorbar_size = colorbar_pos.xmax - colorbar_pos.xmin
+    elif position == "top":
+        pad_distance = colorbar_pos.ymin - fig_pos.ymax
+        colorbar_size = colorbar_pos.ymax - colorbar_pos.ymin
+    elif position == "bottom":
+        pad_distance = fig_pos.ymin - colorbar_pos.ymax
+        colorbar_size = colorbar_pos.ymax - colorbar_pos.ymin
+
+    assert np.isclose(colorbar_size, expected_size)
+    assert np.isclose(pad_distance, expected_pad)
+
+    plt.close(fig1)
+    plt.close(fig2)