diff --git a/qiskit/primitives/containers/observables_array.py b/qiskit/primitives/containers/observables_array.py index 21c415d7589..6ac30e18cd0 100644 --- a/qiskit/primitives/containers/observables_array.py +++ b/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() diff --git a/test/python/primitives/containers/test_observables_array.py b/test/python/primitives/containers/test_observables_array.py index 5a8513a5ed9..e89b6834efe 100644 --- a/test/python/primitives/containers/test_observables_array.py +++ b/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)