Allow highlighting a subset of qubits when displaying heatmap

Yash-10 commented 5 months ago

Issue

For #4691, as part of Unitary Hack 2024.

Checklist

[X] Add implementation to highlight a subset of qubits.
- These are added as options to the plot() method. A subset of qubits is highlighted in the heatmap if a list of qubits is passed; otherwise, the default behavior is retained.
- The implementation mainly involves configuring the collection_options.
[x] Add tests to ensure the implementation is correct and appropriate.
- These test whether the relevant qubits are highlighted. Properties in the collection_options are checked for this purpose.

Example visualizations

A single qubit

import pathlib
import shutil
import string
from tempfile import mkdtemp

import numpy as np
import pytest

import matplotlib as mpl
import matplotlib.pyplot as plt

from cirq.devices import grid_qubit
from cirq.vis import heatmap

value_map = {
    (grid_qubit.GridQubit(0, 0),): 0.1,
    (grid_qubit.GridQubit(0, 1),): 0.2,
    (grid_qubit.GridQubit(0, 2),): 0.3,
    (grid_qubit.GridQubit(1, 0),): 0.4,
}
single_qubit_heatmap = heatmap.Heatmap(value_map)

qubits = np.array([k[0] for k, _ in value_map.items()])
selected_qubits_indices = [1, 3]
selected_qubits = qubits[selected_qubits_indices]

_, ax = plt.subplots()
_ = single_qubit_heatmap.plot(ax, selected_qubits=selected_qubits)

download

Two-qubit interactions

value_map = {
    (grid_qubit.GridQubit(0, 0), grid_qubit.GridQubit(0, 1)): 0.1,
    (grid_qubit.GridQubit(0, 1), grid_qubit.GridQubit(0, 2)): 0.2,
    (grid_qubit.GridQubit(1, 0), grid_qubit.GridQubit(0, 0)): 0.3,
    (grid_qubit.GridQubit(3, 3), grid_qubit.GridQubit(3, 2)): 0.9,
}
two_qubit_interaction_heatmap = heatmap.TwoQubitInteractionHeatmap(value_map)

qubits = np.array(list(value_map.keys()))
selected_qubits_indices = [(0, 1), (2, 1), (3, 0)]
selected_qubits = qubits[tuple(np.transpose(selected_qubits_indices))]

_, ax = plt.subplots()
_ = two_qubit_interaction_heatmap.plot(ax, selected_qubits=selected_qubits)

download

google-cla[bot] commented 5 months ago

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

codecov[bot] commented 5 months ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 97.81%. Comparing base (72f0542) to head (4c9d7aa). Report is 1 commits behind head on main.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #6642 +/- ## ======================================= Coverage 97.81% 97.81% ======================================= Files 1066 1066 Lines 91693 91779 +86 ======================================= + Hits 89687 89773 +86 Misses 2006 2006 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

Yash-10 commented 5 months ago

I am a bit unsure about the expected mypy error:

cirq-core/cirq/vis/heatmap.py:325: error: Item "None" of "Optional[Any]" has no attribute "update"  [union-attr]

because self._config.collection_options, even if not passed, is set to a dictionary by default in __init__ ({"cmap" : "viridis"}), so it should never be None. I understand it's trying to look at all possible types (for example, if, in the future, collection_options is not set to a dictionary), then it may raise this error.

pavoljuhas commented 5 months ago

I am a bit unsure about the expected mypy error:

cirq-core/cirq/vis/heatmap.py:325: error: Item "None" of "Optional[Any]" has no attribute "update"  [union-attr]

That is expected - mypy is a static typechecker and can only analyze possible return types from self._config.get(...) which include None. I would use self._config["collection_options"] instead.

Yash-10 commented 5 months ago

Thank you, @pavoljuhas; the error is resolved now.

Yash-10 commented 5 months ago

The CLA error occurs because I had used one of the commit suggestions. See https://github.com/quantumlib/Cirq/pull/6642/commits/aa1ed053b5f9bb2969093b4df5d84c7ecd92893e.

NoureldinYosri commented 5 months ago

@Yash-10 huh, interesting I didn't notice that github doesn't like my personal account. anyway remove that commit and add the changes yourselve. you may need to the git push -f to clear the commit history of this PR

pavoljuhas commented 5 months ago

@Yash-10 huh, interesting I didn't notice that github doesn't like my personal account. anyway remove that commit and add the changes yourselve. you may need to the git push -f to clear the commit history of this PR

@NoureldinYosri - if the CLA fails only because of your personal email in a commit, we can override it with an explanation note.

NoureldinYosri commented 5 months ago

@pavoljuhas you can override it without any explanation :D

NoureldinYosri commented 5 months ago

~@Yash-10 I think the easiest thing to do is, after you address the Pavol's test comment you can close this PR and open a new clean one with all the changes.~

NVM, pavol knows away to override CLA

pavoljuhas commented 5 months ago

@Yash-10 I think the easiest thing to do is, after you address the Pavol's test comment you can close this PR and open a new clean one with all the changes.

@Yash-10 - sorry about misunderstanding, we should be able to override the missing CLA check on our side (because it is caused by other authors). It is easier to finalize this PR, no need (hopefully) to create a new one.

pavoljuhas commented 5 months ago

Thank you @Yash-10 !

Yash-10 commented 5 months ago

Thank you both of you for your prompt reviews on this!

quantumlib / Cirq