Visualization tools for single-qubit POVMs and products of single-qub…

…it POVMs. (#17)

* add bloch visualization for SQPovm and product of SQPovms

* update type hints and documentation

* add TODO

* add rank-1 checks

* add colorbar option

* add how-to

* fix docs

* add explanations

* not implemented error for `MultiqubitPOVM`

* linting

* documentation

* add explanations to notebook

* update documentation and explanations

* update after review

* add unittest

* fix random seed in notebook

* update docstrings

docs/how_tos/index.rst

@@ -4,8 +4,8 @@ How-To Guides
 
 This page will summarize the available how-to guides.
 
-.. .. toctree::
-..    :maxdepth: 1
-..    :glob:
-..
-..    *
+.. toctree::
+   :maxdepth: 1
+   :glob:
+
+   *

povm_toolbox/quantum_info/multi_qubit_frame.py

@@ -199,13 +199,13 @@ def analysis(
 
     @classmethod
     def from_vectors(cls, frame_vectors: np.ndarray) -> Self:
-        r"""Initialize a frame from non-normalized bloch vectors :math:`|\psi \rangle`.
+        r"""Initialize a frame from non-normalized bloch vectors :math:`|\tilde{\psi} \rangle = \sqrt{\gamma} |\psi \rangle`.
 
         Args:
-            frame_vectors: list of vectors :math:`|\psi \rangle`. The length of the list corresponds to
+            frame_vectors: list of vectors :math:`|\tilde{\psi} \rangle`. The length of the list corresponds to
                 the number of operators of the frame. Each vector is of shape :math:`(\mathrm{dim},)` where :math:`\mathrm{dim}`
                 is the :attr:`.dimension` of the Hilbert space on which the frame acts. The resulting frame
-                operators :math:`\Pi = |\psi \rangle \langle \psi|` are of shape :math:`(\mathrm{dim}, \mathrm{dim})` as expected.
+                operators :math:`\Pi = \gamma |\psi \rangle \langle \psi|` are of shape :math:`(\mathrm{dim}, \mathrm{dim})` as expected.
 
         Returns:
             The frame corresponding to the vectors.

povm_toolbox/quantum_info/multi_qubit_povm.py

@@ -13,6 +13,8 @@
 from __future__ import annotations
 
 import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
 from qiskit.quantum_info import DensityMatrix, Operator, SparsePauliOp, Statevector
 
 from .base_povm import BasePOVM
@@ -72,3 +74,27 @@ def get_prob(
         if not isinstance(rho, SparsePauliOp):
             rho = Operator(rho)
         return self.analysis(rho, outcome_idx)
+
+    def draw_bloch(
+        self,
+        *,
+        title: str = "",
+        figure: Figure | None = None,
+        axes: Axes | list[Axes] | None = None,
+        figsize: tuple[float, float] | None = None,
+        font_size: float | None = None,
+        colorbar: bool = False,
+    ) -> Figure:
+        """Draw the Bloch vectors of a :class:`.MultiQubitPOVM` instance.
+
+        Args:
+            title: A string that represents the plot title.
+            figure: User supplied Matplotlib Figure instance for plotting Bloch sphere.
+            axes: User supplied Matplotlib axes to render the bloch sphere.
+            figsize: Figure size in inches. Has no effect if passing ``ax``.
+            font_size: Size of font used for Bloch sphere labels.
+            colorbar: If ``True``, normalize the vectors on the Bloch sphere and
+                add a colormap to keep track of the norm of the vectors. It can
+                help to visualize the vector if they have a small norm.
+        """
+        raise NotImplementedError

povm_toolbox/quantum_info/product_povm.py

@@ -12,8 +12,12 @@
 
 from __future__ import annotations
 
+import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
 from qiskit.quantum_info import DensityMatrix, SparsePauliOp, Statevector
+from qiskit.visualization.utils import matplotlib_close_if_inline
 
 from .base_povm import BasePOVM
 from .multi_qubit_povm import MultiQubitPOVM
@@ -67,3 +71,70 @@ def get_prob(
         if not isinstance(rho, SparsePauliOp):
             rho = SparsePauliOp.from_operator(rho)
         return self.analysis(rho, outcome_idx)
+
+    def draw_bloch(
+        self,
+        *,
+        title: str = "",
+        figure: Figure | None = None,
+        axes: Axes | list[Axes] | None = None,
+        figsize: tuple[float, float] | None = None,
+        font_size: float | None = None,
+        colorbar: bool = False,
+    ) -> Figure:
+        """Plot a Bloch sphere for each single-qubit POVM forming the product POVM.
+
+        Args:
+            title: a string that represents the plot title.
+            figure: User supplied Matplotlib Figure instance for plotting Bloch sphere.
+            axes: User supplied Matplotlib axes to render the bloch sphere.
+            figsize: size of each individual Bloch sphere figure, in inches.
+            font_size: Font size for the Bloch ball figures.
+            colorbar: If ``True``, normalize the vectors on the Bloch sphere and
+                add a colormap to keep track of the norm of the vectors. It can
+                help to visualize the vector if they have a small norm.
+        """
+        # Number of subplots (one per qubit)
+        num = self.n_subsystems
+
+        # Check that all local POVMs are single-qubit POVMs
+        if any([len(idx) > 1 for idx in self.sub_systems]):
+            raise NotImplementedError
+
+        # Determine the number of rows and columns for the figure
+        n_cols = int(np.sqrt(num) * 4 / 3)
+        n_rows = int(np.sqrt(num) * 3 / 4) or 1
+        while n_cols * n_rows < num:
+            n_cols += 1 if n_cols * n_rows < num else 0
+            n_rows += 1 if n_cols * n_rows < num else 0
+            while (n_cols - 1) * n_rows >= num:
+                n_cols -= 1
+
+        # Set default values
+        if figsize is None:
+            figsize = (5, 4) if colorbar else (5, 5)
+        width, height = figsize
+        width *= n_cols
+        height *= n_rows
+        title_font_size = font_size if font_size is not None else 16
+
+        # Plot figure
+        fig = figure if figure is not None else plt.figure(figsize=(width, height))
+        for i, idx in enumerate(self.sub_systems):
+            ax = (
+                axes[i]
+                if isinstance(axes, list)
+                else fig.add_subplot(n_rows, n_cols, i + 1, projection="3d")
+            )
+            self[idx].draw_bloch(
+                title="qubit " + ", ".join(map(str, idx)),
+                figure=fig,
+                axes=ax,
+                figsize=figsize,
+                font_size=font_size,
+                colorbar=colorbar,
+            )
+        fig.suptitle(title, fontsize=title_font_size, y=1.0)
+        matplotlib_close_if_inline(fig)
+
+        return fig

povm_toolbox/quantum_info/single_qubit_povm.py

@@ -12,6 +12,13 @@
 
 from __future__ import annotations
 
+import matplotlib as mpl
+import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from qiskit.visualization.bloch import Bloch
+from qiskit.visualization.utils import matplotlib_close_if_inline
+
 from .multi_qubit_povm import MultiQubitPOVM
 
 
@@ -26,6 +33,88 @@ def _check_validity(self) -> None:
         """
         if not self.dimension == 2:
             raise ValueError(
-                f"Dimension of Single Qubit POVM operator space should be 2, not {self.dimension}."
+                "Dimension of Single Qubit POVM operator space should be 2,"
+                f" not {self.dimension}."
             )
         super()._check_validity()
+
+    def get_bloch_vectors(self) -> np.ndarray:
+        r"""Compute the Bloch vector of each effect of the POVM.
+
+        For a rank-1 POVM, each effect :math:`M_k` can be written as
+
+        .. math::
+            M_k = \gamma_k |\psi_k \rangle \langle \psi_k | = \gamma_k
+            \frac{1}{2} \left( \mathbb{I} + \vec{a}_k \cdot \vec{\sigma} \right)
+
+        where :math:`\vec{\sigma}` is the usual Pauli vector and
+        :math:`||\vec{a}_k||^2=1`. We then define the Bloch vector of a rank-1
+        effect as :math:`\vec{r}_k = \gamma_k \vec{a}_k`, which uniquely defines
+        the rank-1 effect.
+        """
+        r = np.empty((self.n_outcomes, 3))
+        for i, pauli_op in enumerate(self.pauli_operators):
+            # Check that the povm effect is rank-1:
+            if np.linalg.matrix_rank(self.operators[i]) > 1:
+                raise ValueError(
+                    "Bloch vector is only well-defined for single-qubit rank-1"
+                    f" POVMs. However, the effect number {i} of this POVM has"
+                    f" rank {np.linalg.matrix_rank(self.operators[i])}."
+                )
+            r[i, 0] = 2 * np.real_if_close(pauli_op.get("X", 0))
+            r[i, 1] = 2 * np.real_if_close(pauli_op.get("Y", 0))
+            r[i, 2] = 2 * np.real_if_close(pauli_op.get("Z", 0))
+        return r
+
+    def draw_bloch(
+        self,
+        *,
+        title: str = "",
+        figure: Figure | None = None,
+        axes: Axes | list[Axes] | None = None,
+        figsize: tuple[float, float] | None = None,
+        font_size: float | None = None,
+        colorbar: bool = False,
+    ) -> Figure:
+        """Plot the Bloch vector of each effect of the POVM.
+
+        Args:
+            title: A string that represents the plot title.
+            figure: User supplied Matplotlib Figure instance for plotting Bloch sphere.
+            axes: User supplied Matplotlib axes to render the bloch sphere.
+            figsize: Figure size in inches. Has no effect if passing ``ax``.
+            font_size: Size of font used for Bloch sphere labels.
+            colorbar: If ``True``, normalize the vectors on the Bloch sphere and
+                add a colormap to keep track of the norm of the vectors. It can
+                help to visualize the vector if they have a small norm.
+        """
+        if figsize is None:
+            figsize = (5, 4) if colorbar else (5, 5)
+
+        # Initialize Bloch sphere
+        B = Bloch(fig=figure, axes=axes, font_size=font_size)
+
+        # Compute Bloch vector
+        vectors = self.get_bloch_vectors()
+
+        if colorbar:
+            # Keep track of vector norms through colorbar
+            cmap = mpl.colormaps["viridis"]
+            B.vector_color = [cmap(np.linalg.norm(vec)) for vec in vectors]
+            # Normalize
+            for i in range(len(vectors)):
+                vectors[i] /= np.linalg.norm(vectors[i])
+
+        B.add_vectors(vectors)
+        B.render(title=title)
+
+        if figure is None:
+            figure = B.fig
+            axes = B.axes
+            figure.set_size_inches(figsize[0], figsize[1])
+            matplotlib_close_if_inline(figure)
+
+        if colorbar:
+            figure.colorbar(mpl.cm.ScalarMappable(cmap=cmap), ax=axes, label="weight")
+
+        return figure

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "qiskit>=1.0",
     "qiskit-ibm-runtime>=0.22",
     "qiskit-aer>=0.13",
+    "matplotlib",
 ]
 
 [project.optional-dependencies]
@@ -48,7 +49,6 @@ lint = [
 ]
 notebook-dependencies = [
     "povm_toolbox",
-    "matplotlib",
     "pylatexenc",
 ]
 docs = [
@@ -58,7 +58,6 @@ docs = [
     "sphinx-autodoc-typehints",
     "sphinx-copybutton",
     "nbsphinx",
-    "matplotlib",
 ]
 
 [tool.coverage.run]
@@ -132,7 +131,7 @@ ignore = [
 explicit-preview-rules = true
 
 [tool.ruff.lint.pylint]
-max-args = 6
+max-args = 7
 
 [tool.ruff.lint.extend-per-file-ignores]
 "docs/**/*" = [

test/quantum_info/test_single_qubit_povm.py

@@ -14,6 +14,7 @@
 from unittest import TestCase
 
 import numpy as np
+from povm_toolbox.library import ClassicalShadows, RandomizedProjectiveMeasurements
 from povm_toolbox.quantum_info.single_qubit_povm import SingleQubitPOVM
 from qiskit.quantum_info import Operator
 from scipy.stats import unitary_group
@@ -63,6 +64,69 @@ def test_pauli_decomposition(self):
 
         # also check that the decomposition is correct TODO
 
+    def test_get_bloch_vectors(self):
+        """Test the method :method:`.SingleQubitPOVM.get_bloch_vectors`."""
+
+        sqpovm = SingleQubitPOVM(
+            [
+                0.8 * Operator.from_label("0"),
+                0.8 * Operator.from_label("1"),
+                0.2 * Operator.from_label("+"),
+                0.2 * Operator.from_label("-"),
+            ]
+        )
+        vectors = np.array([[0, 0, 0.8], [0, 0, -0.8], [0.2, 0, 0], [-0.2, 0, 0]])
+        self.assertTrue(np.allclose(sqpovm.get_bloch_vectors(), vectors))
+
+        sqpovm = SingleQubitPOVM(
+            [
+                0.4 * Operator.from_label("-"),
+                0.4 * Operator.from_label("+"),
+                0.6 * Operator.from_label("r"),
+                0.6 * Operator.from_label("l"),
+            ]
+        )
+        vectors = np.array([[-0.4, 0, 0], [0.4, 0, 0], [0, 0.6, 0], [0, -0.6, 0]])
+        self.assertTrue(np.allclose(sqpovm.get_bloch_vectors(), vectors))
+
+        sqpovm = SingleQubitPOVM(
+            [
+                0.4 * Operator(np.eye(2)),
+                0.6 * Operator.from_label("r"),
+                0.6 * Operator.from_label("l"),
+            ]
+        )
+        with self.assertRaises(ValueError):
+            sqpovm.get_bloch_vectors()
+
+        sqpovm = ClassicalShadows(1).definition()[(0,)]
+        vectors = np.array(
+            [
+                [0, 0, 1 / 3],
+                [0, 0, -1 / 3],
+                [1 / 3, 0, 0],
+                [-1 / 3, 0, 0],
+                [0, 1 / 3, 0],
+                [0, -1 / 3, 0],
+            ]
+        )
+        self.assertTrue(np.allclose(sqpovm.get_bloch_vectors(), vectors))
+
+        sqpovm = RandomizedProjectiveMeasurements(
+            1,
+            bias=np.array([0.2, 0.8]),
+            angles=np.array([0.25 * np.pi, 0, 0.25 * np.pi, 0.25 * np.pi]),
+        ).definition()[(0,)]
+        vectors = np.sqrt(0.5) * np.array(
+            [
+                [0.2, 0, 0.2],
+                [-0.2, 0, -0.2],
+                [0.8 * np.sqrt(0.5), 0.8 * np.sqrt(0.5), 0.8],
+                [-0.8 * np.sqrt(0.5), -0.8 * np.sqrt(0.5), -0.8],
+            ]
+        )
+        self.assertTrue(np.allclose(sqpovm.get_bloch_vectors(), vectors))
+
     # TODO: write a unittest for each public method of SingleQubitPOVM
 
     # TODO: write a unittest to assert the correct handling of invalid inputs (i.e. verify that