Sampled expectation value for distributions (#8748)

* Add base code

* Add tests

* Add more tests

* move files to results and move imports

* fix test

* add more tests

* change from dict to counts

* Run black

* Convert sampled exp val function to rust

This commit rewrites the python implementation of the expectation value
calculation to rust. The rust implementation should be significantly
faster to compute the expectation value.

* Fix lint

* ignore cyclic import warning since inside func

* black

* Add some checks on inputs

* Add release note

* Speed up bitstring access in rust code

This commit fixes the primary bottleneck in the rust expectation value
calculation. The bit string access of the nth position previously used
in the code was O(n) and for large numbers of qubits this access could
add up in aggregate while computing the expectation value (even
with profiling overhead for 127 qubit this was still roughly .3
seconds for calling the rust function). This commit adjusts how we
access the nth character of the bit string to be O(1) and removing that
bottlneck and for 127 qubit bit string provides a ~3x speedup. The
tradeoff is that this isn't safe (in that there is potential incorrect
behavior) if someone injects random unicode characters in the bitstring.
However as this is unlikely making the change is worth it considering
the speed up. After this commit the largest bottleneck is conversion
overhead from python->rust and memory management of the temp rust
object, at ~59% of the time to run the sampled_expval_complex()
function, which is hard to workaround. Although, there might be further
tuning to reduce the time required for the remaining 41%.

* Update qiskit/result/sampled_expval.py

If this works then yes indeed.  Each operator appears to be coded differently so happy there is some overlap here

Co-authored-by: Julien Gacon <gaconju@gmail.com>

* explain raw dict inputs

* Update test/python/result/test_sampled_expval.py

Co-authored-by: Julien Gacon <gaconju@gmail.com>

* Apply suggestions from code review

Co-authored-by: John Lapeyre <jlapeyre@users.noreply.github.com>

* Make oper_table_size a const variable instead of a magic number

Co-authored-by: John Lapeyre <jlapeyre@users.noreply.github.com>

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>
Co-authored-by: John Lapeyre <jlapeyre@users.noreply.github.com>
Co-authored-by: Julien Gacon <gaconju@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

qiskit/__init__.py

@@ -31,6 +31,7 @@ sys.modules["qiskit._accelerate.dense_layout"] = qiskit._accelerate.dense_layout
 sys.modules["qiskit._accelerate.sparse_pauli_op"] = qiskit._accelerate.sparse_pauli_op
 sys.modules["qiskit._accelerate.results"] = qiskit._accelerate.results
 sys.modules["qiskit._accelerate.optimize_1q_gates"] = qiskit._accelerate.optimize_1q_gates
+sys.modules["qiskit._accelerate.sampled_exp_val"] = qiskit._accelerate.sampled_exp_val
 
 
 # Extend namespace for backwards compat

qiskit/quantum_info/__init__.py

@@ -163,4 +163,7 @@ from .synthesis import (
     XXDecomposer,
 )
 
-from .analysis import hellinger_distance, hellinger_fidelity
+from .analysis import (
+    hellinger_distance,
+    hellinger_fidelity,
+)

qiskit/quantum_info/analysis/__init__.py

@@ -1,6 +1,6 @@
 # This code is part of Qiskit.
 #
-# (C) Copyright IBM 2017, 2018.
+# (C) Copyright IBM 2017, 2022.
 #
 # This code is licensed under the Apache License, Version 2.0. You may
 # obtain a copy of this license in the LICENSE.txt file in the root directory

qiskit/result/__init__.py

@@ -36,6 +36,14 @@ Distributions
    ProbDistribution
    QuasiDistribution
 
+Expectation values
+==================
+
+.. autosummary::
+   :toctree: ../stubs/
+
+   sampled_expectation_value
+
 Mitigation
 ==========
 .. autosummary::
@@ -54,8 +62,8 @@ from .utils import marginal_distribution
 from .utils import marginal_memory
 from .counts import Counts
 
-from .distributions.probability import ProbDistribution
-from .distributions.quasi import QuasiDistribution
+from .distributions import QuasiDistribution, ProbDistribution
+from .sampled_expval import sampled_expectation_value
 from .mitigation.base_readout_mitigator import BaseReadoutMitigator
 from .mitigation.correlated_readout_mitigator import CorrelatedReadoutMitigator
 from .mitigation.local_readout_mitigator import LocalReadoutMitigator

qiskit/result/distributions/__init__.py

@@ -13,3 +13,5 @@
 """
 Distributions
 """
+from .probability import ProbDistribution
+from .quasi import QuasiDistribution

qiskit/result/sampled_expval.py

@@ -0,0 +1,85 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2022.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+# pylint: disable=cyclic-import
+
+"""Routines for computing expectation values from sampled distributions"""
+import numpy as np
+
+
+# pylint: disable=import-error
+from qiskit._accelerate.sampled_exp_val import sampled_expval_float, sampled_expval_complex
+from qiskit.exceptions import QiskitError
+from .distributions import QuasiDistribution, ProbDistribution
+
+
+# A list of valid diagonal operators
+OPERS = {"Z", "I", "0", "1"}
+
+
+def sampled_expectation_value(dist, oper):
+    """Computes expectation value from a sampled distribution
+
+    Note that passing a raw dict requires bit-string keys.
+
+    Parameters:
+        dist (Counts or QuasiDistribution or ProbDistribution or dict): Input sampled distribution
+        oper (str or Pauli or PauliOp or PauliSumOp or SparsePauliOp): The operator for
+                                                                       the observable
+
+    Returns:
+        float: The expectation value
+    Raises:
+        QiskitError: if the input distribution or operator is an invalid type
+    """
+    from .counts import Counts
+    from qiskit.quantum_info import Pauli, SparsePauliOp
+    from qiskit.opflow import PauliOp, PauliSumOp
+
+    # This should be removed when these return bit-string keys
+    if isinstance(dist, (QuasiDistribution, ProbDistribution)):
+        dist = dist.binary_probabilities()
+
+    if not isinstance(dist, (Counts, dict)):
+        raise QiskitError("Invalid input distribution type")
+    if isinstance(oper, str):
+        oper_strs = [oper.upper()]
+        coeffs = np.asarray([1.0])
+    elif isinstance(oper, Pauli):
+        oper_strs = [oper.to_label()]
+        coeffs = np.asarray([1.0])
+    elif isinstance(oper, PauliOp):
+        oper_strs = [oper.primitive.to_label()]
+        coeffs = np.asarray([1.0])
+    elif isinstance(oper, PauliSumOp):
+        spo = oper.primitive
+        oper_strs = spo.paulis.to_labels()
+        coeffs = np.asarray(spo.coeffs) * oper.coeff
+    elif isinstance(oper, SparsePauliOp):
+        oper_strs = oper.paulis.to_labels()
+        coeffs = np.asarray(oper.coeffs)
+    else:
+        raise QiskitError("Invalid operator type")
+
+    # Do some validation here
+    bitstring_len = len(next(iter(dist)))
+    if any(len(op) != bitstring_len for op in oper_strs):
+        raise QiskitError(
+            f"One or more operators not same length ({bitstring_len}) as input bitstrings"
+        )
+    for op in oper_strs:
+        if set(op).difference(OPERS):
+            raise QiskitError(f"Input operator {op} is not diagonal")
+    # Dispatch to Rust routines
+    if coeffs.dtype == np.dtype(complex).type:
+        return sampled_expval_complex(oper_strs, coeffs, dist)
+    else:
+        return sampled_expval_float(oper_strs, coeffs, dist)

releasenotes/notes/sampled_expval-85e300e0fb5fa5ea.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Adds the ``sampled_expectation_value`` function that allows for computing expectation values
+    for diagonal operators from distributions such as ``Counts`` and ``QuasiDistribution``.
+    Valid operators are: ``str``, ``Pauli``, ``PauliOp``, ``PauliSumOp``, and ``SparsePauliOp``.

src/lib.rs

@@ -23,6 +23,7 @@ mod optimize_1q_gates;
 mod pauli_exp_val;
 mod results;
 mod sabre_swap;
+mod sampled_exp_val;
 mod sparse_pauli_op;
 mod stochastic_swap;
 
@@ -48,5 +49,6 @@ fn _accelerate(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pymodule!(sparse_pauli_op::sparse_pauli_op))?;
     m.add_wrapped(wrap_pymodule!(results::results))?;
     m.add_wrapped(wrap_pymodule!(optimize_1q_gates::optimize_1q_gates))?;
+    m.add_wrapped(wrap_pymodule!(sampled_exp_val::sampled_exp_val))?;
     Ok(())
 }

src/pauli_exp_val.rs

@@ -28,7 +28,7 @@ const PARALLEL_THRESHOLD: usize = 19;
 // https://stackoverflow.com/a/67191480/14033130
 // and adjust for f64 usage
 #[inline]
-fn fast_sum(values: &[f64]) -> f64 {
+pub fn fast_sum(values: &[f64]) -> f64 {
     let chunks = values.chunks_exact(LANES);
     let remainder = chunks.remainder();
 

src/sampled_exp_val.rs

@@ -0,0 +1,94 @@
+// This code is part of Qiskit.
+//
+// (C) Copyright IBM 2022
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use num_complex::Complex64;
+
+use hashbrown::HashMap;
+use numpy::PyReadonlyArray1;
+use pyo3::prelude::*;
+use pyo3::wrap_pyfunction;
+
+use crate::pauli_exp_val::fast_sum;
+
+const OPER_TABLE_SIZE: usize = (b'Z' as usize) + 1;
+const fn generate_oper_table() -> [[f64; 2]; OPER_TABLE_SIZE] {
+    let mut table = [[0.; 2]; OPER_TABLE_SIZE];
+    table[b'Z' as usize] = [1., -1.];
+    table[b'0' as usize] = [1., 0.];
+    table[b'1' as usize] = [0., 1.];
+    table
+}
+
+static OPERS: [[f64; 2]; OPER_TABLE_SIZE] = generate_oper_table();
+
+fn bitstring_expval(dist: &HashMap<String, f64>, mut oper_str: String) -> f64 {
+    let inds: Vec<usize> = oper_str
+        .char_indices()
+        .filter_map(|(index, oper)| if oper != 'I' { Some(index) } else { None })
+        .collect();
+    oper_str.retain(|c| !r#"I"#.contains(c));
+    let denom: f64 = fast_sum(&dist.values().copied().collect::<Vec<f64>>());
+    let exp_val: f64 = dist
+        .iter()
+        .map(|(bits, val)| {
+            let temp_product: f64 = oper_str.bytes().enumerate().fold(1.0, |acc, (idx, oper)| {
+                let diagonal: [f64; 2] = OPERS[oper as usize];
+                let index_char: char = bits.as_bytes()[inds[idx]] as char;
+                let index: usize = index_char.to_digit(10).unwrap() as usize;
+                acc * diagonal[index]
+            });
+            val * temp_product
+        })
+        .sum();
+    exp_val / denom
+}
+
+/// Compute the expectation value from a sampled distribution
+#[pyfunction]
+#[pyo3(text_signature = "(oper_strs, coeff, dist, /)")]
+pub fn sampled_expval_float(
+    oper_strs: Vec<String>,
+    coeff: PyReadonlyArray1<f64>,
+    dist: HashMap<String, f64>,
+) -> PyResult<f64> {
+    let coeff_arr = coeff.as_slice()?;
+    let out = oper_strs
+        .into_iter()
+        .enumerate()
+        .map(|(idx, string)| coeff_arr[idx] * bitstring_expval(&dist, string))
+        .sum();
+    Ok(out)
+}
+
+/// Compute the expectation value from a sampled distribution
+#[pyfunction]
+#[pyo3(text_signature = "(oper_strs, coeff, dist, /)")]
+pub fn sampled_expval_complex(
+    oper_strs: Vec<String>,
+    coeff: PyReadonlyArray1<Complex64>,
+    dist: HashMap<String, f64>,
+) -> PyResult<f64> {
+    let coeff_arr = coeff.as_slice()?;
+    let out: Complex64 = oper_strs
+        .into_iter()
+        .enumerate()
+        .map(|(idx, string)| coeff_arr[idx] * Complex64::new(bitstring_expval(&dist, string), 0.))
+        .sum();
+    Ok(out.re)
+}
+
+#[pymodule]
+pub fn sampled_exp_val(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(sampled_expval_float))?;
+    m.add_wrapped(wrap_pyfunction!(sampled_expval_complex))?;
+    Ok(())
+}

test/python/result/test_sampled_expval.py

@@ -0,0 +1,114 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2017, 2018.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Tests for qiskit.quantum_info.analysis"""
+
+import unittest
+
+from qiskit.result import Counts, QuasiDistribution, ProbDistribution, sampled_expectation_value
+from qiskit.quantum_info import Pauli, SparsePauliOp
+from qiskit.opflow import PauliOp, PauliSumOp
+from qiskit.test import QiskitTestCase
+
+
+PROBS = {
+    "1000": 0.0022,
+    "1001": 0.0045,
+    "1110": 0.0081,
+    "0001": 0.0036,
+    "0010": 0.0319,
+    "0101": 0.001,
+    "1100": 0.0008,
+    "1010": 0.0009,
+    "1111": 0.3951,
+    "0011": 0.0007,
+    "0111": 0.01,
+    "0000": 0.4666,
+    "1101": 0.0355,
+    "1011": 0.0211,
+    "0110": 0.0081,
+    "0100": 0.0099,
+}
+
+
+class TestSampledExpval(QiskitTestCase):
+    """Test sampled expectation values"""
+
+    def test_simple(self):
+        """Test that basic exp values work"""
+
+        dist2 = {"00": 0.5, "11": 0.5}
+        dist3 = {"000": 0.5, "111": 0.5}
+        # ZZ even GHZ is 1.0
+        self.assertAlmostEqual(sampled_expectation_value(dist2, "ZZ"), 1.0)
+        # ZZ odd GHZ is 0.0
+        self.assertAlmostEqual(sampled_expectation_value(dist3, "ZZZ"), 0.0)
+        # All id ops goes to 1.0
+        self.assertAlmostEqual(sampled_expectation_value(dist3, "III"), 1.0)
+        # flipping one to I makes even GHZ 0.0
+        self.assertAlmostEqual(sampled_expectation_value(dist2, "IZ"), 0.0)
+        self.assertAlmostEqual(sampled_expectation_value(dist2, "ZI"), 0.0)
+        # Generic Z on PROBS
+        self.assertAlmostEqual(sampled_expectation_value(PROBS, "ZZZZ"), 0.7554)
+
+    def test_same(self):
+        """Test that all operators agree with each other for counts input"""
+        ans = 0.9356
+        counts = Counts(
+            {
+                "001": 67,
+                "110": 113,
+                "100": 83,
+                "011": 205,
+                "111": 4535,
+                "101": 100,
+                "010": 42,
+                "000": 4855,
+            }
+        )
+        oper = "IZZ"
+
+        exp1 = sampled_expectation_value(counts, oper)
+        self.assertAlmostEqual(exp1, ans)
+
+        exp2 = sampled_expectation_value(counts, Pauli(oper))
+        self.assertAlmostEqual(exp2, ans)
+
+        exp3 = sampled_expectation_value(counts, PauliOp(Pauli(oper)))
+        self.assertAlmostEqual(exp3, ans)
+
+        spo = SparsePauliOp([oper], coeffs=[1])
+        exp4 = sampled_expectation_value(counts, PauliSumOp(spo, coeff=2))
+        self.assertAlmostEqual(exp4, 2 * ans)
+
+        exp5 = sampled_expectation_value(counts, SparsePauliOp.from_list([[oper, 1]]))
+        self.assertAlmostEqual(exp5, ans)
+
+    def test_asym_ops(self):
+        """Test that asymmetric exp values work"""
+        dist = QuasiDistribution(PROBS)
+        self.assertAlmostEqual(sampled_expectation_value(dist, "0III"), 0.5318)
+        self.assertAlmostEqual(sampled_expectation_value(dist, "III0"), 0.5285)
+        self.assertAlmostEqual(sampled_expectation_value(dist, "1011"), 0.0211)
+
+    def test_probdist(self):
+        """Test that ProbDistro"""
+        dist = ProbDistribution(PROBS)
+        result = sampled_expectation_value(dist, "IZIZ")
+        self.assertAlmostEqual(result, 0.8864)
+
+        result2 = sampled_expectation_value(dist, "00ZI")
+        self.assertAlmostEqual(result2, 0.4376)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)