Add ObservablesArray.apply_layout()

qiskit/primitives/containers/observables_array.py

@@ -26,7 +26,9 @@
 import numpy as np
 from numpy.typing import ArrayLike
 
+from qiskit.exceptions import QiskitError
 from qiskit.quantum_info import Pauli, PauliList, SparsePauliOp
+from qiskit.transpiler import TranspileLayout
 
 from .object_array import object_array
 from .shape import ShapedMixin, shape_tuple
@@ -97,6 +99,90 @@ def __repr__(self):
         array = np.array2string(self._array, prefix=prefix, suffix=suffix, threshold=50)
         return prefix + array + suffix
 
+    def apply_layout(
+        self, layout: TranspileLayout | list[int] | None, num_qubits: int | None = None
+    ) -> ObservablesArray:
+        """Apply a transpiler layout to this observables array.
+
+        Args:
+            layout: Either a :class:`~.TranspileLayout`, a list of integers, or ``None``.
+                If both layout and ``num_qubits`` are none, a copy of the operator is
+                returned.
+            num_qubits: The number of qubits to expand the operator to. If not
+                provided then, if ``layout`` is a :class:`~.TranspileLayout`, the
+                number of the transpiler output circuit qubits will be used by
+                default. However, if ``layout`` is a list of integers, the permutation
+                specified will be applied without any expansion. If layout is
+                ``None``, the operator will be expanded to the given ``num_qubits``.
+
+        Returns:
+            A new observables array with the provided layout applied.
+
+        Raises:
+            QiskitError: If a :class:`~.TranspileLayout` is given that maps to qubit numbers that
+                are larger than the number of qubits in this array.
+            QiskitError: If ``num_qubits`` is less than the number of
+        """
+        if layout is None and num_qubits is None or self.size == 0:
+            return ObservablesArray(self._array, copy=True, validate=False)
+
+        # Determine index layout for each observable
+        obs_num_qubits = len(next(iter(self._array.flat[0])))
+
+        if layout is None:
+            n_qubits = obs_num_qubits
+            layout = list(range(n_qubits))
+        elif isinstance(layout, TranspileLayout):
+            n_qubits = len(layout._output_qubit_list)
+            layout = layout.final_index_layout()
+        else:
+            n_qubits = len(layout)
+
+        if num_qubits is not None:
+            if num_qubits < n_qubits:
+                raise QiskitError(
+                    f"The input num_qubits is too small, a {num_qubits} qubit layout cannot be "
+                    f"applied to a {n_qubits} qubit operator"
+                )
+            n_qubits = num_qubits
+        if layout is not None and any(x >= n_qubits for x in layout):
+            raise QiskitError("Provided layout contains indices outside the number of qubits.")
+
+        # Check if layout is trivial mapping
+        trivial_layout = layout == list(range(obs_num_qubits))
+
+        # If trivial layout and no qubit padding we return a copy
+        if trivial_layout and n_qubits == obs_num_qubits:
+            return ObservablesArray(self._array, copy=True, validate=False)
+
+        # Otherwise we need to pad and possible remap all dict keys
+        # This is super inefficient, and we really need a new data
+        # structure to avoid all this iteration and string manipulation
+        if trivial_layout:
+            pad = (n_qubits - obs_num_qubits) * "I"
+
+            def _key_fn(key):
+                return pad + key
+
+        else:
+
+            def _key_fn(key):
+                new_key = n_qubits * ["I"]
+                for char, qubit in zip(reversed(key), layout):
+                    # Qubit position is from end of string
+                    new_key[n_qubits - 1 - qubit] = char
+                return "".join(new_key)
+
+        mapped_array = np.empty_like(self._array)
+        for idx, observable in np.ndenumerate(self._array):
+            # Remap observable
+            new_observable = {}
+            for key, val in observable.items():
+                new_observable[_key_fn(key)] = val
+            mapped_array[idx] = new_observable
+
+        return ObservablesArray(mapped_array, copy=False, validate=False)
+
     def tolist(self) -> list:
         """Convert to a nested list"""
         return self._array.tolist()

test/python/primitives/containers/test_observables_array.py

@@ -12,13 +12,19 @@
 
 """Test ObservablesArray"""
 
+from test import QiskitTestCase
+
 import itertools as it
 import ddt
 import numpy as np
 
+from qiskit.exceptions import QiskitError
+from qiskit.circuit import QuantumCircuit
 import qiskit.quantum_info as qi
+from qiskit.primitives import StatevectorEstimator
 from qiskit.primitives.containers.observables_array import ObservablesArray
-from test import QiskitTestCase  # pylint: disable=wrong-import-order
+from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
+from qiskit.providers.fake_provider import GenericBackendV2
 
 
 @ddt.ddt
@@ -357,3 +363,67 @@ def test_validate(self):
         obs = ObservablesArray([{"XX": 1}, {"XYZ": 1}], validate=False)
         with self.assertRaisesRegex(ValueError, "number of qubits must be the same"):
             obs.validate()
+
+    def test_apply_layout_errors(self):
+        """Test the apply_layout() method errors as expected."""
+        obs = ObservablesArray({"XYZ": 1})
+        layout = [1, 0, 2]
+        with self.assertRaisesRegex(QiskitError, "num_qubits is too small"):
+            obs.apply_layout(layout, 2)
+
+    def test_apply_layout_trivial(self):
+        """Test the apply_layout() method for trivial inputs."""
+        obs = ObservablesArray("XYZ")
+        self.assertEqual(obs.apply_layout(None).tolist(), obs.tolist())
+        self.assertEqual(obs.apply_layout(None, 3).tolist(), obs.tolist())
+
+    def test_apply_layout_padding(self):
+        """Test the apply_layout() method with identity padding."""
+        obs = ObservablesArray({"XYZ": 1, "XYY": 2})
+        self.assertEqual(obs.apply_layout(None, 5).tolist(), {"IIXYZ": 1, "IIXYY": 2})
+        self.assertEqual(obs.apply_layout([2, 1, 0], 5).tolist(), {"IIZYX": 1, "IIYYX": 2})
+        self.assertEqual(obs.apply_layout([4, 1, 0], 5).tolist(), {"ZIIYX": 1, "YIIYX": 2})
+
+    def test_apply_layout_permutations(self):
+        """Test the apply_layout() method under permutations."""
+        obs_in = ObservablesArray([{"XYZ": 1}])
+        obs_out = obs_in.apply_layout([0, 1, 2])
+        self.assertEqual(obs_out.tolist(), [{"XYZ": 1}])
+
+        obs_in = ObservablesArray([{"XYZ": 1}])
+        obs_out = obs_in.apply_layout([1, 0, 2])
+        self.assertEqual(obs_out.tolist(), [{"XZY": 1}])
+
+        obs_in = ObservablesArray([{"XYZ": 1}])
+        obs_out = obs_in.apply_layout([2, 1, 0])
+        self.assertEqual(obs_out.tolist(), [{"ZYX": 1}])
+
+    def test_apply_layout_from_transpiled_circuit(self):
+        """Test that the apply_layout() method interfaces properly with transpiler routing."""
+
+        circuit = QuantumCircuit(5)
+        circuit.h(0)
+        circuit.cx(0, 4)
+        circuit.cx(3, 0)
+        for idx in range(5):
+            circuit.rz(idx / 5, idx)
+        circuit.cx(0, 4)
+        circuit.cx(2, 1)
+        circuit.h(0)
+
+        obs = ObservablesArray(
+            [["I" * idx + pauli + "I" * (4 - idx) for pauli in "XYZ"] for idx in range(5)]
+        )
+
+        backend = GenericBackendV2(num_qubits=5, coupling_map=[[idx, idx + 1] for idx in range(4)])
+        pm = generate_preset_pass_manager(optimization_level=1, backend=backend)
+        isa_circuit = pm.run(circuit)
+        isa_obs = obs.apply_layout(isa_circuit.layout)
+
+        # sanity check that we have a non-trivial layout
+        self.assertNotEqual(isa_circuit.layout.final_index_layout(), [0, 1, 2, 3, 4])
+
+        # the estimates before and after the rerouting should be the same for any observables
+        estimator = StatevectorEstimator()
+        result = estimator.run([(circuit, obs), (isa_circuit, isa_obs)]).result()
+        np.testing.assert_allclose(result[0].data.evs, result[1].data.evs, rtol=0, atol=1e-10)