Implement SparseObservable.apply_layout (#13372)

jakelishman · web-flow · commit 280c4dcb3216 · 2024-10-30T08:59:49.000Z
This is one more usability method to bring `SparseObservable` closer
inline with `SparsePauliOp`. The same functionality is relatively easily
implementable by the user by iterating through the terms, mapping the
indices, and putting the output back into
`SparseObservable.from_sparse_list`, but given how heavily we promote
the method for `SparsePauliOp`, it probably forms part of the core of
user expectations.
diff --git a/crates/accelerate/src/lib.rs b/crates/accelerate/src/lib.rs
@@ -10,6 +10,10 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
+// This stylistic lint suppression should be in `Cargo.toml`, but we can't do that until we're at an
+// MSRV of 1.74 or greater.
+#![allow(clippy::comparison_chain)]
+
 use std::env;
 
 use pyo3::import_exception;
diff --git a/crates/accelerate/src/sparse_observable.rs b/crates/accelerate/src/sparse_observable.rs
@@ -12,6 +12,7 @@
 
 use std::collections::btree_map;
 
+use hashbrown::HashSet;
 use num_complex::Complex64;
 use num_traits::Zero;
 use thiserror::Error;
@@ -263,8 +264,11 @@ impl ::std::convert::TryFrom<u8> for BitTerm {
     }
 }
 
-/// Error cases stemming from data coherence at the point of entry into `SparseObservable` from raw
-/// arrays.
+/// Error cases stemming from data coherence at the point of entry into `SparseObservable` from
+/// user-provided arrays.
+///
+/// These most typically appear during [from_raw_parts], but can also be introduced by various
+/// remapping arithmetic functions.
 ///
 /// These are generally associated with the Python-space `ValueError` because all of the
 /// `TypeError`-related ones are statically forbidden (within Rust) by the language, and conversion
@@ -285,6 +289,10 @@ pub enum CoherenceError {
     DecreasingBoundaries,
     #[error("the values in `indices` are not term-wise increasing")]
     UnsortedIndices,
+    #[error("the input contains duplicate qubits")]
+    DuplicateIndices,
+    #[error("the provided qubit mapping does not account for all contained qubits")]
+    IndexMapTooSmall,
 }
 impl From<CoherenceError> for PyErr {
     fn from(value: CoherenceError) -> PyErr {
@@ -753,7 +761,9 @@ impl SparseObservable {
             let indices = &indices[left..right];
             if !indices.is_empty() {
                 for (index_left, index_right) in indices[..].iter().zip(&indices[1..]) {
-                    if index_left >= index_right {
+                    if index_left == index_right {
+                        return Err(CoherenceError::DuplicateIndices);
+                    } else if index_left > index_right {
                         return Err(CoherenceError::UnsortedIndices);
                     }
                 }
@@ -931,6 +941,42 @@ impl SparseObservable {
         Ok(())
     }
 
+    /// Relabel the `indices` in the operator to new values.
+    ///
+    /// This fails if any of the new indices are too large, or if any mapping would cause a term to
+    /// contain duplicates of the same index.  It may not detect if multiple qubits are mapped to
+    /// the same index, if those qubits never appear together in the same term.  Such a mapping
+    /// would not cause data-coherence problems (the output observable will be valid), but is
+    /// unlikely to be what you intended.
+    ///
+    /// *Panics* if `new_qubits` is not long enough to map every index used in the operator.
+    pub fn relabel_qubits_from_slice(&mut self, new_qubits: &[u32]) -> Result<(), CoherenceError> {
+        for qubit in new_qubits {
+            if *qubit >= self.num_qubits {
+                return Err(CoherenceError::BitIndexTooHigh);
+            }
+        }
+        let mut order = btree_map::BTreeMap::new();
+        for i in 0..self.num_terms() {
+            let start = self.boundaries[i];
+            let end = self.boundaries[i + 1];
+            for j in start..end {
+                order.insert(new_qubits[self.indices[j] as usize], self.bit_terms[j]);
+            }
+            if order.len() != end - start {
+                return Err(CoherenceError::DuplicateIndices);
+            }
+            for (index, dest) in order.keys().zip(&mut self.indices[start..end]) {
+                *dest = *index;
+            }
+            for (bit_term, dest) in order.values().zip(&mut self.bit_terms[start..end]) {
+                *dest = *bit_term;
+            }
+            order.clear();
+        }
+        Ok(())
+    }
+
     /// Return a suitable Python error if two observables do not have equal numbers of qubits.
     fn check_equal_qubits(&self, other: &SparseObservable) -> PyResult<()> {
         if self.num_qubits != other.num_qubits {
@@ -2020,6 +2066,77 @@ impl SparseObservable {
         }
         out
     }
+
+    /// Apply a transpiler layout to this :class:`SparseObservable`.
+    ///
+    /// Typically you will have defined your observable in terms of the virtual qubits of the
+    /// circuits you will use to prepare states.  After transpilation, the virtual qubits are mapped
+    /// to particular physical qubits on a device, which may be wider than your circuit.  That
+    /// mapping can also change over the course of the circuit.  This method transforms the input
+    /// observable on virtual qubits to an observable that is suitable to apply immediately after
+    /// the fully transpiled *physical* circuit.
+    ///
+    /// Args:
+    ///     layout (TranspileLayout | list[int] | None): The layout to apply.  Most uses of this
+    ///         function should pass the :attr:`.QuantumCircuit.layout` field from a circuit that
+    ///         was transpiled for hardware.  In addition, you can pass a list of new qubit indices.
+    ///         If given as explicitly ``None``, no remapping is applied (but you can still use
+    ///         ``num_qubits`` to expand the observable).
+    ///     num_qubits (int | None): The number of qubits to expand the observable to.  If not
+    ///         supplied, the output will be as wide as the given :class:`.TranspileLayout`, or the
+    ///         same width as the input if the ``layout`` is given in another form.
+    ///
+    /// Returns:
+    ///     A new :class:`SparseObservable` with the provided layout applied.
+    #[pyo3(signature = (/, layout, num_qubits=None), name = "apply_layout")]
+    fn py_apply_layout(&self, layout: Bound<PyAny>, num_qubits: Option<u32>) -> PyResult<Self> {
+        let py = layout.py();
+        let check_inferred_qubits = |inferred: u32| -> PyResult<u32> {
+            if inferred < self.num_qubits {
+                return Err(PyValueError::new_err(format!(
+                    "cannot shrink the qubit count in an observable from {} to {}",
+                    self.num_qubits, inferred
+                )));
+            }
+            Ok(inferred)
+        };
+        if layout.is_none() {
+            let mut out = self.clone();
+            out.num_qubits = check_inferred_qubits(num_qubits.unwrap_or(self.num_qubits))?;
+            return Ok(out);
+        }
+        let (num_qubits, layout) = if layout.is_instance(
+            &py.import_bound(intern!(py, "qiskit.transpiler"))?
+                .getattr(intern!(py, "TranspileLayout"))?,
+        )? {
+            (
+                check_inferred_qubits(
+                    layout.getattr(intern!(py, "_output_qubit_list"))?.len()? as u32
+                )?,
+                layout
+                    .call_method0(intern!(py, "final_index_layout"))?
+                    .extract::<Vec<u32>>()?,
+            )
+        } else {
+            (
+                check_inferred_qubits(num_qubits.unwrap_or(self.num_qubits))?,
+                layout.extract()?,
+            )
+        };
+        if layout.len() < self.num_qubits as usize {
+            return Err(CoherenceError::IndexMapTooSmall.into());
+        }
+        if layout.iter().any(|qubit| *qubit >= num_qubits) {
+            return Err(CoherenceError::BitIndexTooHigh.into());
+        }
+        if layout.iter().collect::<HashSet<_>>().len() != layout.len() {
+            return Err(CoherenceError::DuplicateIndices.into());
+        }
+        let mut out = self.clone();
+        out.num_qubits = num_qubits;
+        out.relabel_qubits_from_slice(&layout)?;
+        Ok(out)
+    }
 }
 
 impl ::std::ops::Add<&SparseObservable> for SparseObservable {
diff --git a/test/python/quantum_info/test_sparse_observable.py b/test/python/quantum_info/test_sparse_observable.py
@@ -13,15 +13,19 @@
 # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring
 
 import copy
+import itertools
 import pickle
+import random
 import unittest
 
 import ddt
 import numpy as np
 
-from qiskit.circuit import Parameter
+from qiskit import transpile
+from qiskit.circuit import Measure, Parameter, library, QuantumCircuit
 from qiskit.exceptions import QiskitError
 from qiskit.quantum_info import SparseObservable, SparsePauliOp, Pauli
+from qiskit.transpiler import Target
 
 from test import QiskitTestCase, combine  # pylint: disable=wrong-import-order
 
@@ -39,6 +43,24 @@ def single_cases():
     ]
 
 
+def lnn_target(num_qubits):
+    """Create a simple `Target` object with an arbitrary basis-gate set, and open-path
+    connectivity."""
+    out = Target()
+    out.add_instruction(library.RZGate(Parameter("a")), {(q,): None for q in range(num_qubits)})
+    out.add_instruction(library.SXGate(), {(q,): None for q in range(num_qubits)})
+    out.add_instruction(Measure(), {(q,): None for q in range(num_qubits)})
+    out.add_instruction(
+        library.CXGate(),
+        {
+            pair: None
+            for lower in range(num_qubits - 1)
+            for pair in [(lower, lower + 1), (lower + 1, lower)]
+        },
+    )
+    return out
+
+
 class AllowRightArithmetic:
     """Some type that implements only the right-hand-sided arithmatic operations, and allows
     `SparseObservable` to pass through them.
@@ -1533,3 +1555,158 @@ def test_clear(self, obs):
         num_qubits = obs.num_qubits
         obs.clear()
         self.assertEqual(obs, SparseObservable.zero(num_qubits))
+
+    def test_apply_layout_list(self):
+        self.assertEqual(
+            SparseObservable.zero(5).apply_layout([4, 3, 2, 1, 0]), SparseObservable.zero(5)
+        )
+        self.assertEqual(
+            SparseObservable.zero(3).apply_layout([0, 2, 1], 8), SparseObservable.zero(8)
+        )
+        self.assertEqual(
+            SparseObservable.identity(2).apply_layout([1, 0]), SparseObservable.identity(2)
+        )
+        self.assertEqual(
+            SparseObservable.identity(3).apply_layout([100, 10_000, 3], 100_000_000),
+            SparseObservable.identity(100_000_000),
+        )
+
+        terms = [
+            ("ZYX", (4, 2, 1), 1j),
+            ("", (), -0.5),
+            ("+-rl01", (10, 8, 6, 4, 2, 0), 2.0),
+        ]
+
+        def map_indices(terms, layout):
+            return [
+                (terms, tuple(layout[bit] for bit in bits), coeff) for terms, bits, coeff in terms
+            ]
+
+        identity = list(range(12))
+        self.assertEqual(
+            SparseObservable.from_sparse_list(terms, num_qubits=12).apply_layout(identity),
+            SparseObservable.from_sparse_list(terms, num_qubits=12),
+        )
+        # We've already tested elsewhere that `SparseObservable.from_sparse_list` produces termwise
+        # sorted indices, so these tests also ensure `apply_layout` is maintaining that invariant.
+        backwards = list(range(12))[::-1]
+        self.assertEqual(
+            SparseObservable.from_sparse_list(terms, num_qubits=12).apply_layout(backwards),
+            SparseObservable.from_sparse_list(map_indices(terms, backwards), num_qubits=12),
+        )
+        shuffled = [4, 7, 1, 10, 0, 11, 3, 2, 8, 5, 6, 9]
+        self.assertEqual(
+            SparseObservable.from_sparse_list(terms, num_qubits=12).apply_layout(shuffled),
+            SparseObservable.from_sparse_list(map_indices(terms, shuffled), num_qubits=12),
+        )
+        self.assertEqual(
+            SparseObservable.from_sparse_list(terms, num_qubits=12).apply_layout(shuffled, 100),
+            SparseObservable.from_sparse_list(map_indices(terms, shuffled), num_qubits=100),
+        )
+        expanded = [78, 69, 82, 68, 32, 97, 108, 101, 114, 116, 33]
+        self.assertEqual(
+            SparseObservable.from_sparse_list(terms, num_qubits=11).apply_layout(expanded, 120),
+            SparseObservable.from_sparse_list(map_indices(terms, expanded), num_qubits=120),
+        )
+
+    def test_apply_layout_transpiled(self):
+        base = SparseObservable.from_sparse_list(
+            [
+                ("ZYX", (4, 2, 1), 1j),
+                ("", (), -0.5),
+                ("+-r", (3, 2, 0), 2.0),
+            ],
+            num_qubits=5,
+        )
+
+        qc = QuantumCircuit(5)
+        initial_list = [3, 4, 0, 2, 1]
+        no_routing = transpile(
+            qc, target=lnn_target(5), initial_layout=initial_list, seed_transpiler=2024_10_25_0
+        ).layout
+        # It's easiest here to test against the `list` form, which we verify separately and
+        # explicitly.
+        self.assertEqual(base.apply_layout(no_routing), base.apply_layout(initial_list))
+
+        expanded = transpile(
+            qc, target=lnn_target(100), initial_layout=initial_list, seed_transpiler=2024_10_25_1
+        ).layout
+        self.assertEqual(
+            base.apply_layout(expanded), base.apply_layout(initial_list, num_qubits=100)
+        )
+
+        qc = QuantumCircuit(5)
+        qargs = list(itertools.permutations(range(5), 2))
+        random.Random(2024_10_25_2).shuffle(qargs)
+        for pair in qargs:
+            qc.cx(*pair)
+
+        routed = transpile(qc, target=lnn_target(5), seed_transpiler=2024_10_25_3).layout
+        self.assertEqual(
+            base.apply_layout(routed),
+            base.apply_layout(routed.final_index_layout(filter_ancillas=True)),
+        )
+
+        routed_expanded = transpile(qc, target=lnn_target(20), seed_transpiler=2024_10_25_3).layout
+        self.assertEqual(
+            base.apply_layout(routed_expanded),
+            base.apply_layout(
+                routed_expanded.final_index_layout(filter_ancillas=True), num_qubits=20
+            ),
+        )
+
+    def test_apply_layout_none(self):
+        self.assertEqual(SparseObservable.zero(0).apply_layout(None), SparseObservable.zero(0))
+        self.assertEqual(SparseObservable.zero(0).apply_layout(None, 3), SparseObservable.zero(3))
+        self.assertEqual(SparseObservable.zero(5).apply_layout(None), SparseObservable.zero(5))
+        self.assertEqual(SparseObservable.zero(3).apply_layout(None, 8), SparseObservable.zero(8))
+        self.assertEqual(
+            SparseObservable.identity(0).apply_layout(None), SparseObservable.identity(0)
+        )
+        self.assertEqual(
+            SparseObservable.identity(0).apply_layout(None, 8), SparseObservable.identity(8)
+        )
+        self.assertEqual(
+            SparseObservable.identity(2).apply_layout(None), SparseObservable.identity(2)
+        )
+        self.assertEqual(
+            SparseObservable.identity(3).apply_layout(None, 100_000_000),
+            SparseObservable.identity(100_000_000),
+        )
+
+        terms = [
+            ("ZYX", (2, 1, 0), 1j),
+            ("", (), -0.5),
+            ("+-rl01", (10, 8, 6, 4, 2, 0), 2.0),
+        ]
+        self.assertEqual(
+            SparseObservable.from_sparse_list(terms, num_qubits=12).apply_layout(None),
+            SparseObservable.from_sparse_list(terms, num_qubits=12),
+        )
+        self.assertEqual(
+            SparseObservable.from_sparse_list(terms, num_qubits=12).apply_layout(
+                None, num_qubits=200
+            ),
+            SparseObservable.from_sparse_list(terms, num_qubits=200),
+        )
+
+    def test_apply_layout_failures(self):
+        obs = SparseObservable.from_list([("IIYI", 2.0), ("IIIX", -1j)])
+        with self.assertRaisesRegex(ValueError, "duplicate"):
+            obs.apply_layout([0, 0, 1, 2])
+        with self.assertRaisesRegex(ValueError, "does not account for all contained qubits"):
+            obs.apply_layout([0, 1])
+        with self.assertRaisesRegex(ValueError, "less than the number of qubits"):
+            obs.apply_layout([0, 2, 4, 6])
+        with self.assertRaisesRegex(ValueError, "cannot shrink"):
+            obs.apply_layout([0, 1], num_qubits=2)
+        with self.assertRaisesRegex(ValueError, "cannot shrink"):
+            obs.apply_layout(None, num_qubits=2)
+
+        qc = QuantumCircuit(3)
+        qc.cx(0, 1)
+        qc.cx(1, 2)
+        qc.cx(2, 0)
+        layout = transpile(qc, target=lnn_target(3), seed_transpiler=2024_10_25).layout
+        with self.assertRaisesRegex(ValueError, "cannot shrink"):
+            obs.apply_layout(layout, num_qubits=2)
