Introduce an optional virtual shape for MultiQubitFrame objects (#101)

* introduce virtual shape

* update unittest

* update product frames

* add unittests for multiqubit frames

* update how `shape` is computed in `ProductFrame`

* add unittests for `ProductFrame`

* update docstrings

* add release note

* fix typos

* update after review

* lint and tests

* update after review

povm_toolbox/quantum_info/multi_qubit_frame.py

@@ -24,6 +24,9 @@
 else:
     from typing import override  # pragma: no cover
 
+from math import prod
+from typing import TypeVar
+
 import numpy as np
 from qiskit.exceptions import QiskitError
 from qiskit.quantum_info import Operator, SparsePauliOp
@@ -32,8 +35,14 @@
 
 from .base import BaseFrame
 
+LabelMultiQubitT = TypeVar("LabelMultiQubitT", int, tuple[int, ...])
+"""Each operator in the frame spanning multiple qubits is identified by a label.
+
+This is the type of these labels. They are either integers or tuples of integers.
+"""
 
-class MultiQubitFrame(BaseFrame[int]):
+
+class MultiQubitFrame(BaseFrame[LabelMultiQubitT]):
     """Class that collects all information that any frame of multiple qubits should specify.
 
     This is a representation of an operator-valued vector space frame. The effects are specified as
@@ -45,15 +54,20 @@ class MultiQubitFrame(BaseFrame[int]):
        for more general information.
     """
 
-    def __init__(self, list_operators: list[Operator]) -> None:
+    def __init__(
+        self, list_operators: list[Operator], *, shape: tuple[int, ...] | None = None
+    ) -> None:
         """Initialize from explicit operators.
 
         Args:
             list_operators: list that contains the explicit frame operators. The length of the list
                 is the number of operators of the frame.
+            shape: the shape defining the indexing of operators in ``list_operators``. If ``None``,
+                the default shape is ``(self.num_operators,)``.
 
 
         Raises:
+            ValueError: if the length of ``list_operators`` is not compatible with ``shape``.
             ValueError: if the frame operators do not have a correct shape. They should all be
                 hermitian and of the same dimension.
         """
@@ -65,6 +79,7 @@ def __init__(self, list_operators: list[Operator]) -> None:
         self._informationally_complete: bool
 
         self._num_operators = len(list_operators)
+        self.shape = shape or (self._num_operators,)
         self._dimension = list_operators[0].dim[0]
         for frame_op in list_operators:
             if not (self._dimension == frame_op.dim[0] and self._dimension == frame_op.dim[1]):
@@ -90,7 +105,11 @@ def __init__(self, list_operators: list[Operator]) -> None:
     def __repr__(self) -> str:
         """Return the string representation of a :class:`.MultiQubitFrame` instance."""
         f_subsystems = f"(num_qubits={self.num_subsystems})" if self.num_subsystems > 1 else ""
-        return f"{self.__class__.__name__}{f_subsystems}<{self.num_operators}> at {hex(id(self))}"
+        repr_str = (
+            f"{self.__class__.__name__}{f_subsystems}<{','.join(map(str, self.shape))}> "
+            f"at {hex(id(self))}"
+        )
+        return repr_str
 
     @property
     def informationally_complete(self) -> bool:
@@ -107,6 +126,21 @@ def num_operators(self) -> int:
         """The number of effects of the frame."""
         return self._num_operators
 
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Return the shape of the frame."""
+        return self._shape
+
+    @shape.setter
+    def shape(self, new_shape: tuple[int, ...]) -> None:
+        """Set the shape of the frame."""
+        if prod(new_shape) != self.num_operators:
+            raise ValueError(
+                f"The shape {new_shape} is not compatible with the number of operators in the frame"
+                f" ({self.num_operators})."
+            )
+        self._shape = new_shape
+
     @property
     def operators(self) -> list[Operator]:
         """Return the list of frame operators."""
@@ -148,7 +182,29 @@ def _check_validity(self) -> None:
             if not np.allclose(op, op.adjoint(), atol=1e-5):
                 raise ValueError(f"The {k}-the frame operator is not hermitian.")
 
-    def __getitem__(self, index: slice) -> Operator | list[Operator]:
+    def _ravel_index(self, index: LabelMultiQubitT) -> int:
+        """Ravel a multi-index into a flat index when applicable..
+
+        Args:
+            index: an integer index or multi-index matching the shape of the frame.
+
+        Returns:
+            A flattened integer index.
+
+        Raises:
+            ValueError: if an integer index is supplied for a frame that has a multi-dimensional
+                shape.
+        """
+        if isinstance(index, tuple):
+            return int(np.ravel_multi_index(multi_index=index, dims=self.shape))
+        if len(self.shape) > 1:
+            raise ValueError(
+                f"The integer index `{index}` is invalid because the frame has a {len(self.shape)}-"
+                "dimensional shape."
+            )
+        return index
+
+    def __getitem__(self, index: LabelMultiQubitT) -> Operator | list[Operator]:
         """Return a frame operator or a list of frame operators.
 
         Args:
@@ -157,7 +213,7 @@ def __getitem__(self, index: slice) -> Operator | list[Operator]:
         Returns:
             The operator or list of operators corresponding to the index.
         """
-        return self.operators[index]
+        return self.operators[self._ravel_index(index)]
 
     def __len__(self) -> int:
         """Return the number of operators of the frame."""
@@ -174,17 +230,20 @@ def __array__(self) -> np.ndarray:
     def analysis(
         self,
         hermitian_op: SparsePauliOp | Operator,
-        frame_op_idx: int | set[int] | None = None,
-    ) -> float | dict[int, float] | np.ndarray:
+        frame_op_idx: LabelMultiQubitT | set[LabelMultiQubitT] | None = None,
+    ) -> float | dict[LabelMultiQubitT, float] | np.ndarray:
         if isinstance(hermitian_op, SparsePauliOp):
             hermitian_op = hermitian_op.to_operator()
         op_vectorized = np.conj(matrix_to_double_ket(hermitian_op.data))
 
-        if isinstance(frame_op_idx, int):
-            return float(np.dot(op_vectorized, self._array[:, frame_op_idx]).real)
+        if isinstance(frame_op_idx, (int, tuple)):
+            return float(
+                np.dot(op_vectorized, self._array[:, self._ravel_index(frame_op_idx)]).real
+            )
         if isinstance(frame_op_idx, set):
             return {
-                idx: float(np.dot(op_vectorized, self._array[:, idx]).real) for idx in frame_op_idx
+                idx: float(np.dot(op_vectorized, self._array[:, self._ravel_index(idx)]).real)
+                for idx in frame_op_idx
             }
         if frame_op_idx is None:
             return np.array(np.dot(op_vectorized, self._array).real)

povm_toolbox/quantum_info/product_frame.py

@@ -70,7 +70,6 @@ def __init__(self, frames: dict[tuple[int, ...], T]) -> None:
         subsystem_indices = set()
         self._dimension = 1
         self._num_operators = 1
-        shape: list[int] = []
         for idx, frame in frames.items():
             idx_set = set(idx)
             if len(idx) != len(idx_set):
@@ -94,14 +93,12 @@ def __init__(self, frames: dict[tuple[int, ...], T]) -> None:
             subsystem_indices.update(idx_set)
             self._dimension *= frame.dimension
             self._num_operators *= frame.num_operators
-            shape.append(frame.num_operators)
 
         self._informationally_complete: bool = all(
             [frame.informationally_complete for frame in frames.values()]
         )
 
         self._frames = frames
-        self._shape: tuple[int, ...] = tuple(shape)
 
         self._check_validity()
 
@@ -182,10 +179,15 @@ def num_operators(self) -> int:
         """The number of effects of the frame."""
         return self._num_operators
 
+    @property
+    def _sub_shapes(self) -> tuple[tuple[int, ...], ...]:
+        """Give the shapes of local frames."""
+        return tuple(frame.shape for frame in self._frames.values())
+
     @property
     def shape(self) -> tuple[int, ...]:
-        """Give the number of operators per sub-system."""
-        return self._shape
+        """Give the shape of the product frame."""
+        return tuple(s for shape in self._sub_shapes for s in shape)
 
     @property
     def sub_systems(self) -> list[tuple[int, ...]]:
@@ -201,6 +203,38 @@ def _check_validity(self) -> None:
         for povm in self._frames.values():
             povm._check_validity()
 
+    def _ravel_index(self, index: tuple[int, ...]) -> tuple[int, ...]:
+        """Process a global multi-index into a tuple of flat indices for each local frame.
+
+        Args:
+            index: a large multi-index consisting of local multi-indices, each of those
+                corresponding to a local frame. Therefore, ``index`` is supposed to be a
+                flattened ``tuple[LabelMultiQubitT, ...]``, which is always a ``tuple[int, ...]`.
+
+        Returns:
+            A multi-index consisting of one integer index per local frame. That is, for each
+            sub-system the local multi-index has been raveled.
+
+        Raises:
+            ValueError: if ``index`` does not have the same number of dimensions as the shape of the
+                frame.
+        """
+        if len(index) != len(self.shape):
+            raise ValueError(
+                f"The index {index} does not have the same number of dimensions as the shape of the"
+                f" frame: {self.shape}"
+            )
+
+        index_processed = []
+        start = 0
+        for sub_shape in self._sub_shapes:
+            local_flat_index = np.ravel_multi_index(
+                index[start : start + len(sub_shape)], sub_shape
+            )
+            index_processed.append(int(local_flat_index))
+            start += len(sub_shape)
+        return tuple(index_processed)
+
     def __getitem__(self, sub_system: tuple[int, ...]) -> T:
         r"""Return the :class:`.MultiQubitFrame` acting on the specified sub-system.
 
@@ -239,20 +273,15 @@ def _trace_of_prod(self, operator: SparsePauliOp, frame_op_idx: tuple[int, ...])
         Args:
             operator: the input operator to multiply with a frame operator.
             frame_op_idx: the label specifying the frame operator to use. The frame operator is
-                labeled by a tuple of integers (one index per local frame).
+                labeled by a tuple of integers (possibly multiple integers for one local frame).
 
         Returns:
             The trace of the product of the input operator with the specified frame operator.
-
-        Raises:
-            IndexError: when the provided outcome label (tuple of integers) has a number of integers
-                which does not correspond to the number of local frames making up the product frame.
-            IndexError: when a local index exceeds the number of operators of the corresponding
-                local frame.
-            ValueError: when the output is not a real number.
         """
         p_idx = 0.0 + 0.0j
 
+        index_processed = self._ravel_index(frame_op_idx)
+
         # Second, we iterate over our input operator, ``operator``.
         for label, op_coeff in operator.label_iter():
             summand = op_coeff
@@ -265,29 +294,15 @@ def _trace_of_prod(self, operator: SparsePauliOp, frame_op_idx: tuple[int, ...])
                 # Extract the local Pauli term on the qubit indices of this local POVM.
                 sublabel = "".join(label[-(i + 1)] for i in idx)
                 # Try to obtain the coefficient of the local POVM for this local Pauli term.
+                local_idx = index_processed[j]
                 try:
-                    local_idx = frame_op_idx[j]
                     coeff = povm.pauli_operators[local_idx][sublabel]
                 except KeyError:
                     # If it does not exist, the current summand becomes 0 because it would be
                     # multiplied by 0.
                     summand = 0.0
                     # In this case we can break the iteration over the remaining local POVMs.
                     break
-                except IndexError as exc:
-                    if len(frame_op_idx) <= j:
-                        raise IndexError(
-                            f"The outcome label {frame_op_idx} does not match the expected shape. "
-                            f"It is supposed to contain {len(self._frames)} integers, but has "
-                            f"{len(frame_op_idx)}."
-                        ) from exc
-                    if povm.num_operators <= frame_op_idx[j]:
-                        raise IndexError(
-                            f"Outcome index '{frame_op_idx[j]}' is out of range for the local POVM"
-                            f" acting on subsystems {idx}. This POVM has {povm.num_operators}"
-                            " outcomes."
-                        ) from exc
-                    raise exc
                 else:
                     # If the label does exist, we multiply the coefficient into our summand.
                     # The factor 2^N_qubit comes from Tr[(P_1...P_N)^2] = 2^N.

releasenotes/notes/virtual-shape-multiqubit-frame-adf7d2a13ea1cd56.yaml

@@ -0,0 +1,7 @@
+---
+features:
+  - |
+    The :class:`.MultiQubitFrame` now has a ``shape`` attribute. Before, its shape was implicitly
+    assumed to be ``(len(frame.num_operators),)`` and the operators were indexed by integers. Note
+    that this is is still the default shape and indexing method. However, the frame operators can
+    now also be indexed by multi-indices (tuple of integers) if a custom shape is specified.