Better use for the figsize argument of plot_bloch_multivector #7263
Description
What is the expected enhancement?
The figsize argument of plot_bloch_multivector is currently not used by the visualisation, so it is impossible to change its size (e.g. to shrink it for single-qubit states). The following tiny change to its code would use figsize as a hint for individual bloch ball sizes, enhancing the usability of this visualization:
# Original source code snippet:
...
width, height = plt.figaspect(1 / num)
...
# Modified source code snippet:
...
if figsize is not None:
width, height = figsize
width *= num
else:
width, height = plt.figaspect(1 / num)
...Comparison of a single-qubit bloch sphere visualisation, original (left) vs scaled version available in proposed enhancement (right):
Three-qubit bloch spheres, original (top) vs scaled version available in proposed enhancement (bottom):
Three-qubit bloch spheres, advanced scaled version available in proposed enhancement, leaving additional space between spheres:
Full source code for reference:
def plot_bloch_multivector(state, title="", figsize=None, *, rho=None, reverse_bits=False):
if not HAS_MATPLOTLIB:
raise MissingOptionalLibraryError(
libname="Matplotlib",
name="plot_bloch_multivector",
pip_install="pip install matplotlib",
)
from matplotlib import pyplot as plt
bloch_data = (
_bloch_multivector_data(state)[::-1] if reverse_bits else _bloch_multivector_data(state)
)
num = len(bloch_data)
if figsize is not None:
width, height = figsize
width *= num
else:
width, height = plt.figaspect(1 / num)
fig = plt.figure(figsize=(width, height))
for i in range(num):
pos = num - 1 - i if reverse_bits else i
ax = fig.add_subplot(1, num, i + 1, projection="3d")
plot_bloch_vector(bloch_data[i], "qubit " + str(pos), ax=ax, figsize=figsize)
fig.suptitle(title, fontsize=16)
matplotlib_close_if_inline(fig)
return fig