From 79e78d87dd0d6544f0fc91790c4868388a919e20 Mon Sep 17 00:00:00 2001
From: Jake Lishman <jake.lishman@ibm.com>
Date: Tue, 5 Nov 2024 13:09:14 +0000
Subject: [PATCH] Return `SparseObservable` from sequence indices

---
 crates/accelerate/src/sparse_observable.rs    | 36 ++++++++++++-------
 .../quantum_info/test_sparse_observable.py    |  8 ++---
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/crates/accelerate/src/sparse_observable.rs b/crates/accelerate/src/sparse_observable.rs
index 60f4491e7f2..8cdc94f316f 100644
--- a/crates/accelerate/src/sparse_observable.rs
+++ b/crates/accelerate/src/sparse_observable.rs
@@ -974,7 +974,7 @@ impl SparseObservable {
         out
     }
 
-    /// Get an owned representation of a single sparse term.
+    /// Get a view onto a representation of a single sparse term.
     ///
     /// This is effectively an indexing operation into the [SparseObservable].  Recall that two
     /// [SparseObservable]s that have different term orders can still represent the same object.
@@ -983,15 +983,15 @@ impl SparseObservable {
     /// # Panics
     ///
     /// If the index is out of bounds.
-    pub fn term(&self, index: usize) -> SparseTerm {
+    pub fn term(&self, index: usize) -> SparseTermView {
         debug_assert!(index < self.num_terms(), "index {index} out of bounds");
         let start = self.boundaries[index];
         let end = self.boundaries[index + 1];
-        SparseTerm {
+        SparseTermView {
             num_qubits: self.num_qubits,
             coeff: self.coeffs[index],
-            bit_terms: (&self.bit_terms[start..end]).into(),
-            indices: (&self.indices[start..end]).into(),
+            bit_terms: &self.bit_terms[start..end],
+            indices: &self.indices[start..end],
         }
     }
 
@@ -1314,15 +1314,15 @@ impl SparseObservable {
     }
 
     fn __getitem__(&self, py: Python, index: PySequenceIndex) -> PyResult<Py<PyAny>> {
-        match index.with_len(self.num_terms())? {
-            SequenceIndex::Int(index) => Ok(self.term(index).into_py(py)),
-            indices => Ok(PyList::new_bound(
-                py,
-                indices.iter().map(|index| self.term(index).into_py(py)),
-            )
-            .into_any()
-            .unbind()),
+        let indices = match index.with_len(self.num_terms())? {
+            SequenceIndex::Int(index) => return Ok(self.term(index).to_term().into_py(py)),
+            indices => indices,
+        };
+        let mut out = SparseObservable::zero(self.num_qubits);
+        for index in indices.iter() {
+            out.add_term(self.term(index))?;
         }
+        Ok(out.into_py(py))
     }
 
     fn __repr__(&self) -> String {
@@ -2521,6 +2521,16 @@ pub struct SparseTermView<'a> {
     pub indices: &'a [u32],
 }
 impl<'a> SparseTermView<'a> {
+    /// Convert this `SparseTermView` into an owning [SparseTerm] of the same data.
+    pub fn to_term(&self) -> SparseTerm {
+        SparseTerm {
+            num_qubits: self.num_qubits,
+            coeff: self.coeff,
+            bit_terms: self.bit_terms.into(),
+            indices: self.indices.into(),
+        }
+    }
+
     fn to_sparse_str(self) -> String {
         let coeff = format!("{}", self.coeff).replace('i', "j");
         let paulis = self
diff --git a/test/python/quantum_info/test_sparse_observable.py b/test/python/quantum_info/test_sparse_observable.py
index e27dea4ffc0..dd6fc5800fc 100644
--- a/test/python/quantum_info/test_sparse_observable.py
+++ b/test/python/quantum_info/test_sparse_observable.py
@@ -1840,10 +1840,10 @@ def test_indexing(self):
         ]
         self.assertEqual(obs[0], expected[0])
         self.assertEqual(obs[-2], expected[-2])
-        self.assertEqual(obs[2:4], expected[2:4])
-        self.assertEqual(obs[1::2], expected[1::2])
-        self.assertEqual(obs[:], expected)
-        self.assertEqual(obs[-1:-4:-1], expected[-1:-4:-1])
+        self.assertEqual(obs[2:4], SparseObservable(expected[2:4]))
+        self.assertEqual(obs[1::2], SparseObservable(expected[1::2]))
+        self.assertEqual(obs[:], SparseObservable(expected))
+        self.assertEqual(obs[-1:-4:-1], SparseObservable(expected[-1:-4:-1]))
 
     @ddt.data(
         SparseObservable.identity(0),