Add an option to highlight a subset of qubits in cirq.Heatmap

Is your feature request related to a use case or problem? Please describe.

Describe the solution you’d like Heatmaps are often useful to visualize error metrics for qubits and couplers on the google devices. When selecting a subset of qubits to run an experiment on; it’s useful to consult heatmaps of error metrics. We should add an option to the cirq.Heatmap and cirq.TwoQubitInteractionHeatmap so that we can pass in a list of qubits and highlight it on the heatmap. For example:

This can be done right now by manually specifying a different edge width and color for the selected qubits. For example:

fig, ax = plt.subplots(figsize=(12, 10))
edge_colors = tuple('red' if q in selected_qubits else 'grey' for q in device.qubits)
linestyle = tuple('solid' if q in selected_qubits else 'dashed' for q in device.qubits)
linewidths = tuple(4 if q in selected_qubits else 2 for q in device.qubits)
cirq.Heatmap({q: 0.0 for q in device.qubits}).plot(
    ax=ax,
    collection_options={
        'cmap': 'binary',
        'linewidths': linewidths,
        'edgecolors': edge_colors,
        'linestyles': linestyle,
    },
    plot_colorbar=False,
    annotation_format=None,
)
heatmap._plot_on_axis(ax)
fig.show()

What is the urgency from your perspective for this issue? Is it blocking important work? P1 - I need this no later than the next release (end of quarter)

[tanujkhattar]

@kris524 Yes, the feature is still needed and we would appreciate help!

[kris524]

One last question: Could you suggest a method for PolyCollection or something that could help me get the data in collection_options? This is for testing purposes (I want to compare the expected vs returned contents of linewidths edge_colors linestyle)