Change SparseObservable.canonicalize to HashMap/Vec internals

crates/quantum_info/src/sparse_observable/mod.rs

@@ -12,7 +12,7 @@
 
 mod lookup;
 
-use hashbrown::HashSet;
+use hashbrown::{HashMap, HashSet};
 use indexmap::IndexSet;
 use itertools::Itertools;
 use ndarray::Array2;
@@ -30,6 +30,7 @@ use pyo3::{
     sync::PyOnceLock,
     types::{IntoPyDict, PyList, PyString, PyTuple, PyType},
 };
+use rayon::prelude::*;
 use std::{
     cmp::Ordering,
     collections::btree_map,
@@ -39,6 +40,7 @@ use std::{
 use thiserror::Error;
 
 use qiskit_circuit::{
+    getenv_use_multiple_threads,
     imports::{ImportOnceCell, NUMPY_COPY_ONLY_IF_NEEDED},
     slice::{PySequenceIndex, SequenceIndex},
 };
@@ -615,18 +617,24 @@ impl SparseObservable {
     ///
     /// This function is idempotent.
     pub fn canonicalize(&self, tol: f64) -> SparseObservable {
-        let mut terms = btree_map::BTreeMap::new();
+        let mut terms = HashMap::with_capacity(self.num_terms());
         for term in self.iter() {
             terms
-                .entry((term.indices, term.bit_terms))
+                .entry(term.canonical_key())
                 .and_modify(|c| *c += term.coeff)
                 .or_insert(term.coeff);
         }
+        let mut vec = terms
+            .into_iter()
+            .filter(|(_, coeff)| coeff.norm_sqr() > tol * tol)
+            .collect::<Vec<_>>();
+        if getenv_use_multiple_threads() {
+            vec.par_sort_unstable_by_key(|&(k, _)| k);
+        } else {
+            vec.sort_unstable_by_key(|&(k, _)| k);
+        }
         let mut out = SparseObservable::zero(self.num_qubits);
-        for ((indices, bit_terms), coeff) in terms {
-            if coeff.norm_sqr() <= tol * tol {
-                continue;
-            }
+        for ((indices, bit_terms), coeff) in vec {
             out.coeffs.push(coeff);
             out.bit_terms.extend_from_slice(bit_terms);
             out.indices.extend_from_slice(indices);
@@ -1463,6 +1471,11 @@ impl SparseTermView<'_> {
         format!("({coeff})({paulis})")
     }
 }
+impl<'a> SparseTermView<'a> {
+    fn canonical_key(self) -> (&'a [u32], &'a [BitTerm]) {
+        (self.indices, self.bit_terms)
+    }
+}
 
 /// A mutable view object onto a single term of a [SparseObservable].
 ///
@@ -3227,9 +3240,7 @@ impl PySparseObservable {
     ///         >>> right = SparseObservable.from_list([("XYZ", 1.0/7.0)] * 7)  # doesn't sum to 1.0
     ///         >>> assert left.simplify() != right.simplify()
     ///         >>> assert (left - right).simplify() == SparseObservable.zero(left.num_qubits)
-    #[pyo3(
-        signature = (/, tol=1e-8),
-    )]
+    #[pyo3(signature = (/, tol=1e-8))]
     fn simplify(&self, tol: f64) -> PyResult<Self> {
         let inner = self.inner.read().map_err(|_| InnerReadError)?;
         let simplified = inner.canonicalize(tol);

releasenotes/notes/spo-canonicalize-impl-a06c2220b4ac0435.yaml

@@ -0,0 +1,6 @@
+features_quantum_info:
+  - |
+    The internal algorithm used by :func:`.SparseObservable.canonicalize` changed to use hash-maps
+    and a subsequent sort, instead of on-the-fly B-tree based sorting during the mapping.  The
+    scaling of the algorithm and the results should be unchanged, but the practical performance,
+    especially for very large observables, has improved.