From 81b3f7890a368d9c8a2602ad7742bb252238ceba Mon Sep 17 00:00:00 2001
From: nate stemen <nate@unitary.fund>
Date: Fri, 26 Jan 2024 09:44:45 -0800
Subject: [PATCH 1/9] Refactor classical postprocessing in shadows module
 (#2152)

* remove unused helper functions

* simplify product calculation

* remove unnecessary type ignore comment

* avoid using private _unitary_ method

* import numpy typing library as a whole

this is consistent with how we do it in other
parts of the codebase

* remove unused `fidelity` function

* add more extensive tests

* refactor/simplify `get_single_shot_pauli_fidelity`

* simplify expected result/assertion testing

* refactor `get_pauli_fidelities`

* simplify variable names in `shadow_state_reconstruction`

* distinct variable names

* remove boolean calibration flag

use optional fidelities instead

* simplify tests; use `np.testing` functions for assertions

* better variable naming in `classical_snapshot`

* refactor `expectation_estimation_shadow`

* remove calibration option from docs; add fidelity back

* fix precedence

* remove `print` statements

* compute batch sizes properly; move `-1` into sum

i misunderstood how `np.array_split` worked. it uses the
number of splits, as opposed to the size of each batch.

(-1)^x + (-1)^y =/= (-1)^(x + y)

* remove `use_calibration` option in docs

* add tests for added util functions
---
 docs/source/examples/rshadows_tutorial.md     |   4 +-
 docs/source/examples/shadows_tutorial.md      |   3 -
 docs/source/guide/shadows-1-intro.md          |  74 ++--
 mitiq/shadows/classical_postprocessing.py     | 319 ++++++------------
 mitiq/shadows/shadows.py                      |  16 +-
 mitiq/shadows/shadows_utils.py                | 130 ++++---
 .../test/test_classical_postprocessing.py     | 262 ++++++--------
 mitiq/shadows/test/test_shadows.py            |  10 +-
 mitiq/shadows/test/test_shadows_utils.py      |  60 ++--
 9 files changed, 359 insertions(+), 519 deletions(-)

diff --git a/docs/source/examples/rshadows_tutorial.md b/docs/source/examples/rshadows_tutorial.md
index 246fd51742..b8c8ec6309 100644
--- a/docs/source/examples/rshadows_tutorial.md
+++ b/docs/source/examples/rshadows_tutorial.md
@@ -249,7 +249,7 @@ plt.xticks(
 )
 
 plt.ylabel("Pauli fidelity")
-plt.legend()
+plt.legend();
 ```
 
 
@@ -329,14 +329,12 @@ def compare_shadow_methods(
         
     output_shadow = classical_post_processing(
         shadow_outcomes=shadow_measurement_result,
-        use_calibration=False,
         observables=observables,
         k_shadows=k_shadows,
     )
 
     output_shadow_cal = classical_post_processing(
         shadow_outcomes=shadow_measurement_result,
-        use_calibration=True,
         calibration_results=f_est,
         observables=observables,
         k_shadows=k_shadows,
diff --git a/docs/source/examples/shadows_tutorial.md b/docs/source/examples/shadows_tutorial.md
index 075661e127..f6798371df 100644
--- a/docs/source/examples/shadows_tutorial.md
+++ b/docs/source/examples/shadows_tutorial.md
@@ -179,7 +179,6 @@ shadow_outcomes = shadow_quantum_processing(
 # get shadow reconstruction of the density matrix
 output = classical_post_processing(
     shadow_outcomes,
-    use_calibration=False,
     state_reconstruction=True,
 )
 rho_shadow = output["reconstructed_state"]
@@ -286,7 +285,6 @@ for n_measurement in n_measurement_list:
         # perform shadow state reconstruction
         rho_shadow = classical_post_processing(
             shadow_outcomes=shadow_subset,
-            use_calibration=False,
             state_reconstruction=True,
         )["reconstructed_state"]
 
@@ -474,7 +472,6 @@ for error in epsilon_grid:
     shadow_outputs = shadow_quantum_processing(test_circuits, cirq_executor, r)
     output = classical_post_processing(
         shadow_outcomes=shadow_outputs,
-        use_calibration=False,
         observables=list_of_paulistrings,
         k_shadows=k,
     )
diff --git a/docs/source/guide/shadows-1-intro.md b/docs/source/guide/shadows-1-intro.md
index 8e1daf36ae..83637e057c 100644
--- a/docs/source/guide/shadows-1-intro.md
+++ b/docs/source/guide/shadows-1-intro.md
@@ -10,43 +10,45 @@ kernelspec:
   language: python
   name: python3
 ---
+
 ```{admonition} Note:
-The documentation for Classical Shadows in Mitiq is still under construction. This users guide will change in the future. 
+The documentation for Classical Shadows in Mitiq is still under construction. This users guide will change in the future.
 ```
 
 # How Do I Use Classical Shadows Estimation?
 
-
 The `mitiq.shadows` module facilitates the application of the classical shadows protocol on quantum circuits, designed for tasks like quantum state tomography or expectation value estimation. In addition this module integrates a robust shadow estimation protocol that's tailored to counteract noise. The primary objective of the classical shadow protocol is to extract information from a quantum state using repeated measurements.
 
 The procedure can be broken down as follows:
 
-1. `shadow_quantum_processing`: 
+1. `shadow_quantum_processing`:
+
    - Purpose: Execute quantum processing on the provided quantum circuit.
    - Outcome: Measurement results from the processed circuit.
 
-2. `classical_post_processing`: 
+2. `classical_post_processing`:
    - Purpose: Handle classical processing of the measurement results.
    - Outcome: Estimation based on user-defined inputs.
 
 For users aiming to employ the robust shadow estimation protocol, an initial step is needed which entails characterizing the noisy quantum channel. This is done by:
 
 0. `pauli_twirling_calibration`
+
    - Purpose: Characterize the noisy quantum channel.
    - Outcome: A dictionary of `calibration_results`.
 
 1. `shadow_quantum_processing`: same as above.
 
 2. `classical_post_processing`
-   - Args: `use_calibration` = True, 
-           `calibration_results` = output of `pauli_twirling_calibration`
+   - Args: `calibration_results` = output of `pauli_twirling_calibration`
    - Outcome: Error mitigated estimation based on user-defined inputs.
 
 Notes:
-   - The calibration process is specifically designed to mitigate noise encountered during the classical shadow protocol, such as rotation and computational basis measurements. It does not address noise that occurs during state preparation.
-   - Do not need to redo the calibration stage (0. `pauli_twirling_calibration`) if:
-       1. The input circuit has a consistent number of qubits.
-       2. The estimated observables have the same or fewer qubit support.
+
+- The calibration process is specifically designed to mitigate noise encountered during the classical shadow protocol, such as rotation and computational basis measurements. It does not address noise that occurs during state preparation.
+- Do not need to redo the calibration stage (0. `pauli_twirling_calibration`) if:
+  1.  The input circuit has a consistent number of qubits.
+  2.  The estimated observables have the same or fewer qubit support.
 
 ## Protocol Overview
 
@@ -54,8 +56,9 @@ The classical shadow protocol aims to create an approximate classical representa
 
 One can use the `mitiq.shadows' module as follows.
 
-### User-defined inputs 
-Define a quantum circuit, e.g., a circuit which prepares a GHZ state with $n$ = `3` qubits, 
+### User-defined inputs
+
+Define a quantum circuit, e.g., a circuit which prepares a GHZ state with $n$ = `3` qubits,
 
 ```{code-cell} ipython3
 import numpy as np
@@ -76,13 +79,12 @@ circuit = cirq.Circuit(
 print(circuit)
 ```
 
-Define an executor to run the circuit on a quantum computer or a noisy simulator. Note that the *robust shadow estimation* technique can only calibrate and mitigate the noise acting on the operations associated to the classical shadow protocol. So, in order to test the technique, we assume that the state preparation part of the circuit is noiseless. In particular, we define an executor in which:
+Define an executor to run the circuit on a quantum computer or a noisy simulator. Note that the _robust shadow estimation_ technique can only calibrate and mitigate the noise acting on the operations associated to the classical shadow protocol. So, in order to test the technique, we assume that the state preparation part of the circuit is noiseless. In particular, we define an executor in which:
 
 1. A noise channel is added to circuit right before the measurements. I.e. $U_{\Lambda_U}(M_z)_{\Lambda_{\mathcal{M}_Z}}\equiv U\Lambda\mathcal{M}_Z$.
 
 2. A single measurement shot is taken for each circuit, as required by classical shadow protocol.
 
-
 ```{code-cell} ipython3
 from mitiq import MeasurementResult
 
@@ -125,8 +127,7 @@ def cirq_executor(
     return executor
 ```
 
-Given the above general executor, we define a specific example of a noisy executor, assuming a bit flip channel with a  probability of `0.1'
-
+Given the above general executor, we define a specific example of a noisy executor, assuming a bit flip channel with a probability of `0.1'
 
 ```{code-cell} ipython3
 from functools import partial
@@ -164,24 +165,28 @@ f_est
 
 the varible `locality` is the maximum number of qubits on which our operators of interest are acting on. E.g. if our operator is a sequence of two point correlation terms $\{\langle Z_iZ_{i+1}\rangle\}_{0\leq i\leq n-1}$, then `locality` = 2. We note that one could also split the calibration process into two stages:
 
-01. `shadow_quantum_processing`
-   - Outcome: Get quantum measurement result of the calibration circuit $|0\rangle^{\otimes n}$ `zero_state_shadow_outcomes`.
+1.  `shadow_quantum_processing`
 
-02. `pauli_twirling_calibration`
-   - Outcome: A dictionary of `calibration_results`.
-For more details, please refer to  [this tutorial](../examples/rshadows_tutorial.md)
+- Outcome: Get quantum measurement result of the calibration circuit $|0\rangle^{\otimes n}$ `zero_state_shadow_outcomes`.
+
+2.  `pauli_twirling_calibration`
+
+- Outcome: A dictionary of `calibration_results`.
+  For more details, please refer to [this tutorial](../examples/rshadows_tutorial.md)
 
 ### 1. Quantum Processing
+
 In this step, we obtain classical shadow snapshots from the input state (before applying the invert channel).
 
 #### 1.1 Add Rotation Gate and Meausure the Rotated State in Computational Basis
+
 At present, the implementation supports random Pauli measurement. This is equivalent to randomly sampling $U$ from the local Clifford group $Cl_2^n$, followed by a $Z$-basis measurement (see [this tutorial](../examples/shadows_tutorial.md) for a clear explanation).
 
 #### 1.2 Get the Classical Shadows
-One can obtain the list of measurement results of local Pauli measurements in terms of bitstrings, and the related Pauli-basis measured in terms of strings as follows.
 
+One can obtain the list of measurement results of local Pauli measurements in terms of bitstrings, and the related Pauli-basis measured in terms of strings as follows.
 
-You have two choices: run the quantum measurement or directly use the results from the previous run. 
+You have two choices: run the quantum measurement or directly use the results from the previous run.
 
 - If **True**, the measurement will be run again.
 - If **False**, the results from the previous run will be used.
@@ -206,20 +211,20 @@ else:
         num_total_measurements_shadow=5000,
     )
 ```
-                                                                     
 
 As an example, we print out one of those measurement outcomes and the associated measured operator:
 
-
 ```{code-cell} ipython3
 print("one snapshot measurement result = ", shadow_measurement_output[0][0])
 print("one snapshot measurement basis = ", shadow_measurement_output[1][0])
 ```
 
 ### 2. Classical Post-Processing
-In this step, we estimate our object of interest (expectation value or density matrix) by post-processing the (previously obtained) measurement outcomes. 
+
+In this step, we estimate our object of interest (expectation value or density matrix) by post-processing the (previously obtained) measurement outcomes.
 
 #### 2.1 Example: Operator Expectation Value Esitimation
+
 For example, if we want to estimate the two point correlation function $\{\langle Z_iZ_{i+1}\rangle\}_{0\leq i\leq n-1}$, we will define the corresponding Puali strings:
 
 ```{code-cell} ipython3
@@ -238,13 +243,11 @@ The corresponding expectation values can be estimated (with and without calibrat
 ```{code-cell} ipython3
 est_corrs = classical_post_processing(
     shadow_outcomes=shadow_measurement_output,
-    use_calibration=False,
     observables=two_pt_correlations,
     k_shadows=1,
 )
 cal_est_corrs = classical_post_processing(
     shadow_outcomes=shadow_measurement_output,
-    use_calibration=True,
     calibration_results=f_est,
     observables=two_pt_correlations,
     k_shadows=1,
@@ -253,7 +256,6 @@ cal_est_corrs = classical_post_processing(
 
 Let's compare the results with the exact theoretical values:
 
-
 ```{code-cell} ipython3
 expval_exact = []
 state_vector = circuit.final_state_vector()
@@ -264,7 +266,6 @@ for i, pauli_string in enumerate(two_pt_correlations):
     expval_exact.append(exp.real)
 ```
 
-
 ```{code-cell} ipython3
 print("Classical shadow estimation:", est_corrs)
 print("Robust shadow estimation   :", cal_est_corrs)
@@ -277,12 +278,10 @@ print(
 )
 ```
 
-
 #### 2.2 Example: GHZ State Reconstruction
-In addition to the estimation of expectation values, the `mitiq.shadow` module can also be used to reconstruct an approximated version of the density matrix. 
-As an example, we use the 3-qubit GHZ circuit, previously defined. As a first step, we calculate the Pauli fidelities $f_b$ characterizing the noisy quantum channel $\mathcal{M}=\sum_{b\in\{0,1\}^n}f_b\Pi_b$:
-
 
+In addition to the estimation of expectation values, the `mitiq.shadow` module can also be used to reconstruct an approximated version of the density matrix.
+As an example, we use the 3-qubit GHZ circuit, previously defined. As a first step, we calculate the Pauli fidelities $f_b$ characterizing the noisy quantum channel $\mathcal{M}=\sum_{b\in\{0,1\}^n}f_b\Pi_b$:
 
 ```{code-cell} ipython3
 noisy_executor = partial(
@@ -307,9 +306,7 @@ else:
 f_est
 ```
 
-
-Similar to the previous case (estimation of expectation values), the quantum processing for estimating the density matrix  is done as follows.
-
+Similar to the previous case (estimation of expectation values), the quantum processing for estimating the density matrix is done as follows.
 
 ```{code-cell} ipython3
 if not run_quantum_processing:
@@ -328,12 +325,10 @@ else:
 ```{code-cell} ipython3
 est_corrs = classical_post_processing(
     shadow_outcomes=shadow_measurement_output,
-    use_calibration=False,
     state_reconstruction=True,
 )
 cal_est_corrs = classical_post_processing(
     shadow_outcomes=shadow_measurement_output,
-    use_calibration=True,
     calibration_results=f_est,
     state_reconstruction=True,
 )
@@ -349,7 +344,6 @@ ghz_true = ghz_state @ ghz_state.conj().T
 ptm_ghz_state = operator_ptm_vector_rep(ghz_true)
 ```
 
-
 ```{code-cell} ipython3
 from mitiq.shadows.shadows_utils import fidelity
 
diff --git a/mitiq/shadows/classical_postprocessing.py b/mitiq/shadows/classical_postprocessing.py
index 5b6cd6247f..64b9ace4a4 100644
--- a/mitiq/shadows/classical_postprocessing.py
+++ b/mitiq/shadows/classical_postprocessing.py
@@ -8,50 +8,38 @@
 # LICENSE file in the root directory of this source tree.
 """Classical post-processing process of classical shadows."""
 
-from itertools import combinations
+from collections import defaultdict
+from functools import reduce
+from itertools import compress
+from operator import mul
+from statistics import median
 from typing import Any, Dict, List, Optional, Tuple
 
 import cirq
 import numpy as np
-from numpy.typing import NDArray
+import numpy.typing as npt
 
 import mitiq
-from mitiq.shadows.shadows_utils import create_string
+from mitiq.shadows.shadows_utils import (
+    batch_calibration_data,
+    create_string,
+    valid_bitstrings,
+)
 from mitiq.utils import matrix_kronecker_product, operator_ptm_vector_rep
 
 # Local unitaries to measure Pauli operators in the Z basis
 PAULI_MAP = {
-    "X": cirq.H._unitary_(),
-    "Y": cirq.H._unitary_() @ cirq.S._unitary_().conj(),
-    "Z": cirq.I._unitary_(),
+    "X": cirq.unitary(cirq.H),
+    "Y": cirq.unitary(cirq.H) @ cirq.unitary(cirq.S).conj(),
+    "Z": cirq.unitary(cirq.I),
 }
 
-# Density matrices of single-qubit basis states
-ZERO_STATE = np.diag([1.0 + 0.0j, 0.0 + 0.0j])
-ONE_STATE = np.diag([0.0 + 0.0j, 1.0 + 0.0j])
-
-# F_LOCAL_MAP is based on local Pauli fidelity of qubit i
-# f_b_i = <<b_i|U_i|P_z^b_i>>
-# s.t. f_0 = U_11U_11^* + U_12U_12^*, f_1 = U_21U_21^* + U_22U_22^*
-F_LOCAL_MAP = {
-    "0X": 0.0,
-    "0Y": 0.0,
-    "0Z": 1.0,
-    "1X": 0.0,
-    "1Y": 0.0,
-    "1Z": -1.0,
-}
-
-"""
-The following functions are used in the classical post-processing
-of calibration
-"""
+ZERO_STATE = np.array([[1, 0], [0, 0]])
+ONE_STATE = np.array([[0, 0], [0, 1]])
 
 
 def get_single_shot_pauli_fidelity(
-    bit_string: str,
-    pauli_string: str,
-    locality: Optional[int] = None,
+    bitstring: str, paulistring: str, locality: Optional[int] = None
 ) -> Dict[str, float]:
     r"""
     Calculate Pauli fidelity :math:`f_b` for a single shot measurement of the
@@ -84,31 +72,23 @@ def get_single_shot_pauli_fidelity(
         than or equal to w. The corresponding Pauli fidelity is the product of
         local Pauli fidelity where the associated locus in the keys are '1'.
     """
-    num_qubits = len(bit_string)
-    if locality is None:
-        locality = num_qubits
-    # local_pauli_fidelity is a list of local Pauli fidelity for each qubit
-    local_pauli_fidelity = np.array(
-        [F_LOCAL_MAP[b + u] for b, u in zip(bit_string, pauli_string)]
-    )
-    # f_est is a dictionary of Pauli fidelity for each b_string
-    f_est = {create_string(num_qubits, []): 1.0}
-    for w in range(1, locality + 1):
-        target_locs = np.array(list(combinations(range(num_qubits), w)))
-        single_round_pauli_fidelity = np.prod(
-            local_pauli_fidelity[target_locs], axis=1
-        )
-        for loc, fidelity in zip(target_locs, single_round_pauli_fidelity):
-            # b_str is a string of length n with maximum number of w 1s.
-            b_str = create_string(num_qubits, loc)
-            f_est[b_str] = fidelity
+    pauli_fidelity = {"Z0": 1.0, "Z1": -1.0}
+    local_fidelities = [
+        pauli_fidelity.get(p + b, 0.0) for b, p in zip(bitstring, paulistring)
+    ]
+    num_qubits = len(bitstring)
+    bitstrings = valid_bitstrings(num_qubits, max_hamming_weight=locality)
+    fidelities = {}
+    for bitstring in bitstrings:
+        subset_fidelities = compress(local_fidelities, map(int, bitstring))
+        fidelities[bitstring] = reduce(mul, subset_fidelities, 1.0)
 
-    return f_est
+    return fidelities
 
 
 def get_pauli_fidelities(
-    calibration_measurement_outcomes: Tuple[List[str], List[str]],
-    k_calibration: int,
+    calibration_outcomes: Tuple[List[str], List[str]],
+    num_batches: int,
     locality: Optional[int] = None,
 ) -> Dict[str, complex]:
     r"""
@@ -120,7 +100,7 @@ def get_pauli_fidelities(
     Args:
         calibration_measurement_outcomes: The `random_Pauli_measurement`
             outcomes for the state :math:`|0\rangle^{\otimes n}`}` .
-        k_calibration: number of splits in the median of means estimator.
+        num_batches: The number of batches in the median of means estimator.
         locality: The locality of the operator, whose expectation value is
             going to be estimated by the classical shadow. E.g., if the
             operator is the Ising model Hamiltonian with nearest neighbor
@@ -130,75 +110,40 @@ def get_pauli_fidelities(
         A :math:`2^n`-dimensional dictionary of Pauli fidelities
         :math:`f_b` for :math:`b = \{0,1\}^{n}`
     """
-
-    # classical values of random Pauli measurement stored in classical computer
-    b_lists, u_lists = calibration_measurement_outcomes
-
-    # number of measurements in each split
-    n_total_measurements = len(b_lists)
-
-    means: Dict[str, List[float]] = {}  # key is bitstring, value is mean
-
-    group_idxes = np.array_split(
-        np.arange(n_total_measurements), k_calibration
-    )
-    # loop over the splits of the shadow:
-    for idxes in group_idxes:
-        b_lists_k = np.array(b_lists)[idxes]
-        u_lists_k = np.array(u_lists)[idxes]
-
-        n_group_measurements = len(b_lists_k)
-        group_results = []
-        for j in range(n_group_measurements):
-            bitstring, u_list = b_lists_k[j], u_lists_k[j]
-            f_est = get_single_shot_pauli_fidelity(
-                bitstring, u_list, locality=locality
+    means = defaultdict(list)
+    for bitstrings, paulistrings in batch_calibration_data(
+        calibration_outcomes, num_batches
+    ):
+        all_fidelities = defaultdict(list)
+        for bitstring, paulistring in zip(bitstrings, paulistrings):
+            fidelities = get_single_shot_pauli_fidelity(
+                bitstring, paulistring, locality=locality
             )
-            group_results.append(f_est)
-
-        f_est_group = {
-            b: sum([f[b] for f in group_results]) / n_group_measurements
-            for b in group_results[0]
-        }
-
-        for bitstring, mean in f_est_group.items():
-            if bitstring not in means:
-                means[bitstring] = []
-            means[bitstring].append(mean)
-    # return the median of means
-    medians = {
-        bitstring: complex(np.median(values))
-        for bitstring, values in means.items()
-    }
-    return medians
+            for b, f in fidelities.items():
+                all_fidelities[b].append(f)
 
+        for bitstring, fids in all_fidelities.items():
+            means[bitstring].append(sum(fids) / num_batches)
 
-"""
-The following functions are used in the classical post-processing of
-classical shadows.
-"""
+    return {
+        bitstring: median(averages) for bitstring, averages in means.items()
+    }
 
 
 def classical_snapshot(
-    b_list_shadow: str,
-    u_list_shadow: str,
-    pauli_twirling_calibration: bool,
-    f_est: Optional[Dict[str, float]] = None,
-) -> NDArray[Any]:
+    bitstring: str,
+    paulistring: str,
+    fidelities: Optional[Dict[str, float]] = None,
+) -> npt.NDArray[Any]:
     r"""
     Implement a single snapshot state reconstruction
     with calibration of the noisy quantum channel.
 
     Args:
-        b_list_shadow: The list of length 1, classical outcomes for the
-            snapshot. Here,
-            b = '0' corresponds to :math:`|0\rangle`, and
-            b = '1' corresponds to :math:`|1\rangle`.
-        u_list_shadow: list of len 1, contains str of ("XYZ..") for
-            the applied Pauli measurement on each qubit.
-        pauli_twirling_calibration: Whether to use Pauli twirling
-            calibration.
-        f_est: The estimated Pauli fidelity for each calibration
+        bitstring: A bitstring corresponding to the outcome ... TODO
+        paulistring: String of the applied Pauli measurement on each qubit.
+        f_est: The estimated Pauli fidelities to use for calibration if
+            available.
 
     Returns:
         Reconstructed classical snapshot in terms of nparray.
@@ -206,48 +151,36 @@ def classical_snapshot(
     # calibrate the noisy quantum channel, output in PTM rep.
     # ptm rep of identity
     I_ptm = operator_ptm_vector_rep(np.eye(2) / np.sqrt(2))
-    # define projections Pi_0 and Pi_1
     pi_zero = np.outer(I_ptm, I_ptm)
     pi_one = np.eye(4) - pi_zero
     pi_zero = np.diag(pi_zero)
     pi_one = np.diag(pi_one)
 
-    if pauli_twirling_calibration:
-        if f_est is None:
-            raise ValueError(
-                "estimation of Pauli fidelity must be provided for Pauli"
-                "twirling calibration."
-            )
+    if fidelities:
         elements = []
-        # get b_list and f for each calibration measurement
-        for b_list_cal, f in f_est.items():
-            pi_snapshot_vecter = []
-            for b_1, b2, u2 in zip(b_list_cal, b_list_shadow, u_list_shadow):
+        for bits, fidelity in fidelities.items():
+            pi_snapshot_vector = []
+            for b1, b2, pauli in zip(bits, bitstring, paulistring):
                 # get pi for each qubit based on calibration measurement
-                pi = pi_zero if b_1 == "0" else pi_one
+                pi = pi_zero if b1 == "0" else pi_one
                 # get state for each qubit based on shadow measurement
                 state = ZERO_STATE if b2 == "0" else ONE_STATE
-                # get U2 for each qubit based on shadow measurement
-                U2 = PAULI_MAP[u2]
-                pi_snapshot_vecter.append(
-                    pi * operator_ptm_vector_rep(U2.conj().T @ state @ U2)
+                # get U for each qubit based on shadow measurement
+                U = PAULI_MAP[pauli]
+                pi_snapshot_vector.append(
+                    pi * operator_ptm_vector_rep(U.conj().T @ state @ U)
                 )
-                # solve for the snapshot state
             elements.append(
-                1 / f * matrix_kronecker_product(pi_snapshot_vecter)
+                1 / fidelity * matrix_kronecker_product(pi_snapshot_vector)
             )
-        rho_snapshot_vector = np.sum(elements, axis=0)
-        # normalize the snapshot state
-        rho_snapshot = rho_snapshot_vector  # * normalize_factor
-    # w/o calibration, noted here, the output in terms of matrix,
-    # not in PTM rep.
+        rho_snapshot = np.sum(elements, axis=0)
     else:
         local_rhos = []
-        for b, u in zip(b_list_shadow, u_list_shadow):
-            state = ZERO_STATE if b == "0" else ONE_STATE
-            U = PAULI_MAP[u]
+        for bit, pauli in zip(bitstring, paulistring):
+            state = ZERO_STATE if bit == "0" else ONE_STATE
+            U = PAULI_MAP[pauli]
             # apply inverse of the quantum channel,get PTM vector rep
-            local_rho = 3.0 * (U.conj().T @ state @ U) - cirq.I._unitary_()
+            local_rho = 3.0 * (U.conj().T @ state @ U) - cirq.unitary(cirq.I)
             local_rhos.append(local_rho)
 
         rho_snapshot = matrix_kronecker_product(local_rhos)
@@ -256,36 +189,25 @@ def classical_snapshot(
 
 def shadow_state_reconstruction(
     shadow_measurement_outcomes: Tuple[List[str], List[str]],
-    pauli_twirling_calibration: bool,
-    f_est: Optional[Dict[str, float]] = None,
-) -> NDArray[Any]:
+    fidelities: Optional[Dict[str, float]] = None,
+) -> npt.NDArray[Any]:
     """Reconstruct a state approximation as an average over all snapshots.
 
     Args:
         shadow_measurement_outcomes: Measurement result and the basis
             performing the measurement obtained from `random_pauli_measurement`
             for classical shadow protocol.
-        shadow_measurement_outcomes: Measurement results obtained from
-            `random_pauli_measurement` for classical shadow protocol.
-        pauli_twirling_calibration: Whether to use Pauli twirling
-            calibration.
-        f_est: The estimated Pauli fidelity for each calibration
+        f_est: The estimated Pauli fidelities to use for calibration if
+            available.
     Returns:
         The state reconstructed from classical shadow protocol
     """
+    bitstrings, paulistrings = shadow_measurement_outcomes
 
-    # classical values of random Pauli measurement stored in classical computer
-    b_lists_shadow, u_lists_shadow = shadow_measurement_outcomes
-
-    # Averaging over snapshot states.
     return np.mean(
         [
-            classical_snapshot(
-                b_list_shadow, u_list_shadow, pauli_twirling_calibration, f_est
-            )
-            for b_list_shadow, u_list_shadow in zip(
-                b_lists_shadow, u_lists_shadow
-            )
+            classical_snapshot(bitstring, paulistring, fidelities)
+            for bitstring, paulistring in zip(bitstrings, paulistrings)
         ],
         axis=0,
     )
@@ -293,10 +215,9 @@ def shadow_state_reconstruction(
 
 def expectation_estimation_shadow(
     measurement_outcomes: Tuple[List[str], List[str]],
-    pauli_str: mitiq.PauliString,  # type: ignore
-    k_shadows: int,
-    pauli_twirling_calibration: bool,
-    f_est: Optional[Dict[str, float]] = None,
+    pauli: mitiq.PauliString,
+    num_batches: int,
+    fidelities: Optional[Dict[str, float]] = None,
 ) -> float:
     """Calculate the expectation value of an observable from classical shadows.
     Use median of means to ameliorate the effects of outliers.
@@ -306,72 +227,42 @@ def expectation_estimation_shadow(
             `z_basis_measurement`.
         pauli_str: Single mitiq observable consisting of
             Pauli operators.
-        k_shadows: number of splits in the median of means estimator.
-        pauli_twirling_calibration: Whether to use Pauli twirling
-            calibration.
-        f_est: The estimated Pauli fidelities for each calibration
+        num_batches: Number of batches to process measurement outcomes in.
+        f_est: The estimated Pauli fidelities to use for calibration if
+            available.
 
     Returns:
-        Float corresponding to the estimate of the observable
-        expectation value.
+        Float corresponding to the estimate of the observable expectation
+        value.
     """
-    num_qubits = len(measurement_outcomes[0][0])
-    obs = pauli_str._pauli
-    coeff = pauli_str.coeff
-
-    target_obs, target_locs = [], []
-    for qubit, pauli in obs.items():
-        target_obs.append(str(pauli))
-        target_locs.append(int(qubit))
-
-    # classical values stored in classical computer
-    b_lists_shadow = np.array([list(u) for u in measurement_outcomes[0]])[
-        :, target_locs
+    bitstrings, paulistrings = measurement_outcomes
+    num_qubits = len(bitstrings[0])
+
+    qubits = sorted(pauli.support())
+    filtered_bitstrings = [
+        "".join([bitstring[q] for q in qubits]) for bitstring in bitstrings
     ]
-    u_lists_shadow = np.array([list(u) for u in measurement_outcomes[1]])[
-        :, target_locs
+    filtered_paulis = [
+        "".join([pauli[q] for q in qubits]) for pauli in paulistrings
     ]
+    filtered_data = (filtered_bitstrings, filtered_paulis)
 
     means = []
-
-    # loop over the splits of the shadow:
-    group_idxes = np.array_split(np.arange(len(b_lists_shadow)), k_shadows)
-
-    # loop over the splits of the shadow:
-    for idxes in group_idxes:
-        matching_indexes = np.nonzero(
-            np.all(u_lists_shadow[idxes] == target_obs, axis=1)
-        )
-
-        if len(matching_indexes[0]):
-            product = (-1) ** np.sum(
-                b_lists_shadow[idxes][matching_indexes].astype(int),
-                axis=1,
-            )
-
-            if pauli_twirling_calibration:
-                if f_est is None:
-                    raise ValueError(
-                        "estimation of Pauli fidelity must be provided for"
-                        "Pauli twirling calibration."
-                    )
-
-                b = create_string(num_qubits, target_locs)
-                f_val = f_est.get(b, None)
-                if f_val is None:
-                    product = 0.0
-                else:
-                    # product becomes an array of snapshots expectation values
-                    # witch satisfy condition (1) and (2)
-                    product = (1.0 / f_val) * product
+    for bits, paulis in batch_calibration_data(filtered_data, num_batches):
+        matching_indices = [i for i, p in enumerate(paulis) if p == pauli.spec]
+        if matching_indices:
+            matching_bits = (bits[i] for i in matching_indices)
+            product = sum((-1) ** bit.count("1") for bit in matching_bits)
+
+            if fidelities:
+                b = create_string(num_qubits, qubits)
+                product /= fidelities.get(b, np.inf)
             else:
-                product = 3 ** (len(target_locs)) * product
+                product *= 3 ** len(qubits)
 
         else:
             product = 0.0
 
-        # append the mean of the product in each split
-        means.append(np.sum(product) / len(idxes))
+        means.append(product / len(bits))
 
-    # return the median of means
-    return float(np.real(np.median(means) * coeff))
+    return np.real(np.median(means) * pauli.coeff)
diff --git a/mitiq/shadows/shadows.py b/mitiq/shadows/shadows.py
index da9f13b86b..9d33779bf4 100644
--- a/mitiq/shadows/shadows.py
+++ b/mitiq/shadows/shadows.py
@@ -154,7 +154,6 @@ def shadow_quantum_processing(
 
 def classical_post_processing(
     shadow_outcomes: Tuple[List[str], List[str]],
-    use_calibration: bool = False,
     calibration_results: Optional[Dict[str, float]] = None,
     observables: Optional[List[mitiq.PauliString]] = None,
     k_shadows: Optional[int] = None,
@@ -166,7 +165,6 @@ def classical_post_processing(
 
     Args:
         shadow_outcomes: The output of function `shadow_quantum_processing`.
-        use_calibration: Whether to use the robust shadow estimation.
         calibration_results: The output of function `pauli_twirling_calibrate`.
         observables: The set of observables to measure.
         k_shadows: Number of groups of "median of means" used for shadow
@@ -175,6 +173,7 @@ def classical_post_processing(
             the expectation value of the observables.
 
     Returns:
+        TODO: rewrite this.
         If `state_reconstruction` is True: state tomography matrix in
         :math:`\mathbb{M}_{2^n}(\mathbb{C})` if use_calibration is False,
         otherwise state tomography vector in :math:`\mathbb{C}^{4^d}`.
@@ -182,12 +181,6 @@ def classical_post_processing(
         observables.
     """
 
-    if use_calibration:
-        if calibration_results is None:
-            raise ValueError(
-                "Calibration results cannot be None when use_calibration"
-            )
-
     """
     Additional information:
     Shadow stage 2: Estimate the expectation value of the observables OR
@@ -196,7 +189,7 @@ def classical_post_processing(
     output: Dict[str, Union[float, NDArray[Any]]] = {}
     if state_reconstruction:
         reconstructed_state = shadow_state_reconstruction(
-            shadow_outcomes, use_calibration, f_est=calibration_results
+            shadow_outcomes, fidelities=calibration_results
         )
         output["reconstructed_state"] = reconstructed_state  # type: ignore
     elif observables is not None:
@@ -207,9 +200,8 @@ def classical_post_processing(
             expectation_values = expectation_estimation_shadow(
                 shadow_outcomes,
                 obs,
-                k_shadows=k_shadows,
-                pauli_twirling_calibration=use_calibration,
-                f_est=calibration_results,
+                num_batches=k_shadows,
+                fidelities=calibration_results,
             )
             output[str(obs)] = expectation_values
     return output
diff --git a/mitiq/shadows/shadows_utils.py b/mitiq/shadows/shadows_utils.py
index a4250a5e6f..8600db6830 100644
--- a/mitiq/shadows/shadows_utils.py
+++ b/mitiq/shadows/shadows_utils.py
@@ -9,41 +9,15 @@
 
 """Defines utility functions for classical shadows protocol."""
 
-from typing import Iterable, List, Tuple
+from typing import Generator, List, Optional, Tuple
 
 import numpy as np
-from numpy.typing import NDArray
+import numpy.typing as npt
 from scipy.linalg import sqrtm
 
 import mitiq
 
 
-def eigenvalues_to_bitstring(values: Iterable[int]) -> str:
-    """Converts eigenvalues to bitstring. e.g., ``[-1,1,1] -> "100"``
-
-    Args:
-        values: A list of eigenvalues (must be $-1$ and $1$).
-
-    Returns:
-        A string of 1s and 0s corresponding to the states associated to
-        eigenvalues.
-    """
-    return "".join(["1" if v == -1 else "0" for v in values])
-
-
-def bitstring_to_eigenvalues(bitstring: str) -> List[int]:
-    """Converts bitstring to eigenvalues. e.g., ``"100" -> [-1,1,1]``
-
-    Args:
-        bitstring: A string of 1s and 0s.
-
-    Returns:
-        A list of eigenvalues (either $-1$ or $1$) corresponding to the
-        bitstring.
-    """
-    return [1 if b == "0" else -1 for b in bitstring]
-
-
 def create_string(str_len: int, loc_list: List[int]) -> str:
     """
     This function returns a string of length ``str_len`` with 1s at the
@@ -67,6 +41,80 @@ def create_string(str_len: int, loc_list: List[int]) -> str:
     )
 
 
+def valid_bitstrings(
+    num_qubits: int, max_hamming_weight: Optional[int] = None
+) -> set[str]:
+    """
+    Description.
+
+    Args:
+        num_qubits:
+        max_hamming_weight:
+
+    Returns:
+        The set of all valid bitstrings on ``num_qubits`` bits, with a maximum
+        hamming weight.
+    Raises:
+        Value error when ``max_hamming_weight`` is not greater than 0.
+    """
+    if max_hamming_weight and max_hamming_weight < 1:
+        raise ValueError(
+            "max_hamming_weight must be greater than 0. "
+            f"Got {max_hamming_weight}."
+        )
+
+    bitstrings = {
+        bin(i)[2:].zfill(num_qubits)
+        for i in range(2**num_qubits)
+        if bin(i).count("1") <= (max_hamming_weight or num_qubits)
+    }
+    return bitstrings
+
+
+def fidelity(
+    sigma: npt.NDArray[np.complex64], rho: npt.NDArray[np.complex64]
+) -> float:
+    """
+    Calculate the fidelity between two states.
+
+    Args:
+        sigma: A state in terms of square matrix or vector.
+        rho: A state in terms square matrix or vector.
+
+    Returns:
+        Scalar corresponding to the fidelity.
+    """
+    if sigma.ndim == 1 and rho.ndim == 1:
+        val = np.abs(np.dot(sigma.conj(), rho)) ** 2.0
+    elif sigma.ndim == 1 and rho.ndim == 2:
+        val = np.abs(sigma.conj().T @ rho @ sigma)
+    elif sigma.ndim == 2 and rho.ndim == 1:
+        val = np.abs(rho.conj().T @ sigma @ rho)
+    elif sigma.ndim == 2 and rho.ndim == 2:
+        val = np.abs(np.trace(sqrtm(sigma) @ rho @ sqrtm(sigma)))
+    else:
+        raise ValueError("Invalid input dimensions")
+    return float(val)
+
+
+def batch_calibration_data(
+    data: Tuple[List[str], List[str]], num_batches: int
+) -> Generator[Tuple[List[str], List[str]], None, None]:
+    """Batch calibration into chunks of size batch_size.
+
+    Args:
+        data: The random Pauli measurement outcomes.
+        batch_size: Size of each batch that will be processed.
+
+    Yields:
+        Tuples of bit strings and pauli strings.
+    """
+    bits, paulis = data
+    batch_size = len(bits) // num_batches
+    for i in range(0, len(bits), batch_size):
+        yield bits[i : i + batch_size], paulis[i : i + batch_size]
+
+
 def n_measurements_tomography_bound(epsilon: float, num_qubits: int) -> int:
     """
     This function returns the minimum number of classical shadows required
@@ -135,29 +183,3 @@ def n_measurements_opts_expectation_bound(
         / error**2
     )
     return int(np.ceil(N * K)), int(K)
-
-
-def fidelity(
-    sigma: NDArray[np.complex64], rho: NDArray[np.complex64]
-) -> float:
-    """
-    Calculate the fidelity between two states.
-
-    Args:
-        sigma: A state in terms of square matrix or vector.
-        rho: A state in terms square matrix or vector.
-
-    Returns:
-        Scalar corresponding to the fidelity.
-    """
-    if sigma.ndim == 1 and rho.ndim == 1:
-        val = np.abs(np.dot(sigma.conj(), rho)) ** 2.0
-    elif sigma.ndim == 1 and rho.ndim == 2:
-        val = np.abs(sigma.conj().T @ rho @ sigma)
-    elif sigma.ndim == 2 and rho.ndim == 1:
-        val = np.abs(rho.conj().T @ sigma @ rho)
-    elif sigma.ndim == 2 and rho.ndim == 2:
-        val = np.abs(np.trace(sqrtm(sigma) @ rho @ sqrtm(sigma)))
-    else:
-        raise ValueError("Invalid input dimensions")
-    return float(val)
diff --git a/mitiq/shadows/test/test_classical_postprocessing.py b/mitiq/shadows/test/test_classical_postprocessing.py
index 9d08b5be6f..46d66d2778 100644
--- a/mitiq/shadows/test/test_classical_postprocessing.py
+++ b/mitiq/shadows/test/test_classical_postprocessing.py
@@ -5,7 +5,6 @@
 
 """Tests for classical post-processing functions for classical shadows."""
 
-import cirq
 import numpy as np
 
 import mitiq
@@ -24,6 +23,65 @@ def test_get_single_shot_pauli_fidelity():
     u_list = "XY"
     expected_result = {"00": 1.0, "01": 0.0, "10": 0.0, "11": 0.0}
     assert get_single_shot_pauli_fidelity(b_list, u_list) == expected_result
+    b_list = "01101"
+    u_list = "XYZYZ"
+    assert get_single_shot_pauli_fidelity(b_list, u_list) == {
+        "00000": 1.0,
+        "10000": 0.0,
+        "01000": 0.0,
+        "00100": -1.0,
+        "00010": 0.0,
+        "00001": -1.0,
+        "11000": 0.0,
+        "10100": 0.0,
+        "10010": 0.0,
+        "10001": 0.0,
+        "01100": 0.0,
+        "01010": 0.0,
+        "01001": 0.0,
+        "00110": 0.0,
+        "00101": 1.0,
+        "00011": 0.0,
+        "11100": 0.0,
+        "11010": 0.0,
+        "11001": 0.0,
+        "10110": 0.0,
+        "10101": 0.0,
+        "10011": 0.0,
+        "01110": 0.0,
+        "01101": 0.0,
+        "01011": 0.0,
+        "00111": 0.0,
+        "11110": 0.0,
+        "11101": 0.0,
+        "11011": 0.0,
+        "10111": 0.0,
+        "01111": 0.0,
+        "11111": 0.0,
+    }
+
+
+def test_get_single_shot_pauli_fidelity_with_locality():
+    b_list = "11101"
+    u_list = "XYZYZ"
+    assert get_single_shot_pauli_fidelity(b_list, u_list, locality=2) == {
+        "00000": 1.0,
+        "10000": 0.0,
+        "01000": 0.0,
+        "00100": -1.0,
+        "00010": 0.0,
+        "00001": -1.0,
+        "11000": 0.0,
+        "10100": 0.0,
+        "10010": 0.0,
+        "10001": 0.0,
+        "01100": 0.0,
+        "01010": 0.0,
+        "01001": 0.0,
+        "00110": 0.0,
+        "00101": 1.0,
+        "00011": 0.0,
+    }
 
 
 def test_get_pauli_fidelity():
@@ -32,24 +90,16 @@ def test_get_pauli_fidelity():
         ["XX", "YY", "ZZ", "XY"],
     )
     k_calibration = 2
-    expected_result = {
-        "00": (1 + 0j),
-        "10": (-0.25 + 0j),
-        "01": (0.25 + 0j),
-        "11": (-0.25 + 0j),
-    }
+    expected_result = {"00": 1, "10": -0.25, "01": 0.25, "11": -0.25}
     result = get_pauli_fidelities(
         calibration_measurement_outcomes, k_calibration
     )
-
-    for key in expected_result.keys():
-        assert np.isclose(result[key], expected_result[key])
+    assert result == expected_result
 
 
 def test_classical_snapshot_cal():
     b_list_shadow = "01"
     u_list_shadow = "XY"
-    pauli_twirling_calibration = True
     f_est = {"00": 1, "01": 1 / 3, "10": 1 / 3, "11": 1 / 9}
     expected_result = operator_ptm_vector_rep(
         np.array(
@@ -62,9 +112,7 @@ def test_classical_snapshot_cal():
         )
     )
     np.testing.assert_array_almost_equal(
-        classical_snapshot(
-            b_list_shadow, u_list_shadow, pauli_twirling_calibration, f_est
-        ),
+        classical_snapshot(b_list_shadow, u_list_shadow, f_est),
         expected_result,
     )
 
@@ -74,166 +122,73 @@ def test_classical_snapshot():
     u_list = "XY"
     expected_result = np.array(
         [
-            [0.25 + 0.0j, 0.0 + 0.75j, 0.75 + 0.0j, 0.0 + 2.25j],
-            [0.0 - 0.75j, 0.25 + 0.0j, 0.0 - 2.25j, 0.75 + 0.0j],
-            [0.75 + 0.0j, 0.0 + 2.25j, 0.25 + 0.0j, 0.0 + 0.75j],
-            [0.0 - 2.25j, 0.75 + 0.0j, 0.0 - 0.75j, 0.25 + 0.0j],
+            [0.25, 0.75j, 0.75, 2.25j],
+            [-0.75j, 0.25, -2.25j, 0.75],
+            [0.75, 2.25j, 0.25, 0.75j],
+            [-2.25j, 0.75, -0.75j, 0.25],
         ]
     )
     result = classical_snapshot(b_list, u_list, False)
-    assert isinstance(result, np.ndarray)
-    assert result.shape == (
-        2 ** len(b_list),
-        2 ** len(b_list),
-    )
-    assert np.allclose(result, expected_result)
+    np.testing.assert_allclose(result, expected_result)
 
 
 def test_shadow_state_reconstruction():
-    b_lists = ["010", "001", "000"]
-    u_lists = ["XYZ", "ZYX", "YXZ"]
-
-    measurement_outcomes = (b_lists, u_lists)
+    bitstrings = ["010", "001", "000"]
+    paulistrings = ["XYZ", "ZYX", "YXZ"]
+    measurement_outcomes = (bitstrings, paulistrings)
 
-    expected_result = np.array(
+    expected_state = np.array(
         [
-            [
-                [
-                    0.5 + 0.0j,
-                    -0.5 + 0.0j,
-                    0.5 + 0.0j,
-                    0.0 + 1.5j,
-                    0.5 - 0.5j,
-                    0.0 + 0.0j,
-                    0.0 + 0.0j,
-                    0.0 + 0.0j,
-                ],
-                [
-                    -0.5 + 0.0j,
-                    0.0 + 0.0j,
-                    0.0 + 1.5j,
-                    -0.25 - 0.75j,
-                    0.0 + 0.0j,
-                    -0.25 + 0.25j,
-                    0.0 + 0.0j,
-                    0.0 + 0.0j,
-                ],
-                [
-                    0.5 + 0.0j,
-                    0.0 - 1.5j,
-                    0.5 + 0.0j,
-                    -0.5 + 0.0j,
-                    0.0 - 3.0j,
-                    0.0 + 0.0j,
-                    0.5 - 0.5j,
-                    0.0 + 0.0j,
-                ],
-                [
-                    0.0 - 1.5j,
-                    -0.25 + 0.75j,
-                    -0.5 + 0.0j,
-                    0.0 + 0.0j,
-                    0.0 + 0.0j,
-                    0.0 + 1.5j,
-                    0.0 + 0.0j,
-                    -0.25 + 0.25j,
-                ],
-                [
-                    0.5 + 0.5j,
-                    0.0 + 0.0j,
-                    0.0 + 3.0j,
-                    0.0 + 0.0j,
-                    0.25 + 0.0j,
-                    0.25 + 0.0j,
-                    0.5 + 0.75j,
-                    0.0 - 0.75j,
-                ],
-                [
-                    0.0 + 0.0j,
-                    -0.25 - 0.25j,
-                    0.0 + 0.0j,
-                    0.0 - 1.5j,
-                    0.25 + 0.0j,
-                    -0.25 + 0.0j,
-                    0.0 - 0.75j,
-                    -0.25 + 0.0j,
-                ],
-                [
-                    0.0 + 0.0j,
-                    0.0 + 0.0j,
-                    0.5 + 0.5j,
-                    0.0 + 0.0j,
-                    0.5 - 0.75j,
-                    0.0 + 0.75j,
-                    0.25 + 0.0j,
-                    0.25 + 0.0j,
-                ],
-                [
-                    0.0 + 0.0j,
-                    0.0 + 0.0j,
-                    0.0 + 0.0j,
-                    -0.25 - 0.25j,
-                    0.0 + 0.75j,
-                    -0.25 + 0.0j,
-                    0.25 + 0.0j,
-                    -0.25 + 0.0j,
-                ],
-            ]
+            [0.5, -0.5, 0.5, 1.5j, 0.5 - 0.5j, 0, 0, 0],
+            [-0.5, 0, 1.5j, -0.25 - 0.75j, 0, -0.25 + 0.25j, 0, 0],
+            [0.5, -1.5j, 0.5, -0.5, -3.0j, 0, 0.5 - 0.5j, 0],
+            [-1.5j, -0.25 + 0.75j, -0.5, 0, 0, 1.5j, 0, -0.25 + 0.25j],
+            [0.5 + 0.5j, 0, 3.0j, 0, 0.25, 0.25, 0.5 + 0.75j, -0.75j],
+            [0, -0.25 - 0.25j, 0, -1.5j, 0.25, -0.25, -0.75j, -0.25],
+            [0, 0, 0.5 + 0.5j, 0, 0.5 - 0.75j, 0.75j, 0.25, 0.25],
+            [0, 0, 0, -0.25 - 0.25j, 0.75j, -0.25, 0.25, -0.25],
         ]
     )
 
-    result = shadow_state_reconstruction(measurement_outcomes, False)
-    num_qubits = len(measurement_outcomes[0])
-    assert isinstance(result, np.ndarray)
-    assert result.shape == (
-        2**num_qubits,
-        2**num_qubits,
-    )
-    assert np.allclose(result, expected_result)
+    state = shadow_state_reconstruction(measurement_outcomes)
+    np.testing.assert_almost_equal(state, expected_state)
 
 
 def test_shadow_state_reconstruction_cal():
-    b_lists = ["01", "01"]
-    u_lists = ["XY", "XY"]
-    measurement_outcomes = (b_lists, u_lists)
-    f_est = {"00": 1, "01": 1 / 3, "10": 1 / 3, "11": 1 / 9}
-    expected_result_vec = operator_ptm_vector_rep(
+    bitstrings, paulistrings = ["01", "01"], ["XY", "XY"]
+    measurement_outcomes = (bitstrings, paulistrings)
+    fidelities = {"00": 1, "01": 1 / 3, "10": 1 / 3, "11": 1 / 9}
+
+    expected_state_vec = operator_ptm_vector_rep(
         np.array(
             [
-                [0.25 + 0.0j, 0.0 + 0.75j, 0.75 + 0.0j, 0.0 + 2.25j],
-                [0.0 - 0.75j, 0.25 + 0.0j, 0.0 - 2.25j, 0.75 + 0.0j],
-                [0.75 + 0.0j, 0.0 + 2.25j, 0.25 + 0.0j, 0.0 + 0.75j],
-                [0.0 - 2.25j, 0.75 + 0.0j, 0.0 - 0.75j, 0.25 + 0.0j],
+                [0.25, 0.75j, 0.75, 2.25j],
+                [-0.75j, 0.25, -2.25j, 0.75],
+                [0.75, 2.25j, 0.25, 0.75j],
+                [-2.25j, 0.75, -0.75j, 0.25],
             ]
         )
     )
-    result = shadow_state_reconstruction(measurement_outcomes, True, f_est)
-    num_qubits = len(measurement_outcomes[0])
-    assert isinstance(result, np.ndarray)
-    assert result.shape == (4**num_qubits,)
-    assert np.allclose(result, expected_result_vec)
+    state = shadow_state_reconstruction(measurement_outcomes, fidelities)
+    np.testing.assert_almost_equal(state, expected_state_vec)
 
 
 def test_expectation_estimation_shadow():
-    b_lists = ["0101", "0110"]
-    u_lists = ["ZZXX", "ZZXX"]
-
-    measurement_outcomes = (b_lists, u_lists)
-    observable = mitiq.PauliString("ZZ", support=(0, 1))
-    k = 1
+    measurement_outcomes = ["0101", "0110"], ["ZZXX", "ZZXX"]
+    pauli = mitiq.PauliString("ZZ")
+    batch_size = 1
     expected_result = -9
 
     result = expectation_estimation_shadow(
-        measurement_outcomes, observable, k, False
+        measurement_outcomes, pauli, batch_size
     )
-    assert isinstance(result, float), f"Expected a float, got {type(result)}"
     assert np.isclose(result, expected_result)
 
 
 def test_expectation_estimation_shadow_cal():
-    b_lists = ["0101", "0110"]
-    u_lists = ["YXZZ", "XXXX"]
-    f_est = {
+    bitstrings = ["0101", "0110"]
+    paulistrings = ["YXZZ", "XXXX"]
+    fidelities = {
         "0000": 1,
         "0001": 1 / 3,
         "0010": 1 / 3,
@@ -252,16 +207,14 @@ def test_expectation_estimation_shadow_cal():
         "1111": 1 / 81,
     }
 
-    measurement_outcomes = b_lists, u_lists
-    observable = mitiq.PauliString("YXZZ", support=(0, 1, 2, 3))
-    k = 1
+    measurement_outcomes = bitstrings, paulistrings
+    pauli = mitiq.PauliString("YXZZ")
+    batch_size = 1
     expected_result = 81 / 2
-    print("expected_result", expected_result)
 
     result = expectation_estimation_shadow(
-        measurement_outcomes, observable, k, True, f_est
+        measurement_outcomes, pauli, batch_size, fidelities
     )
-    assert isinstance(result, float), f"Expected a float, got {type(result)}"
     assert np.isclose(result, expected_result)
 
 
@@ -270,13 +223,12 @@ def test_expectation_estimation_shadow_no_indices():
     Test expectation estimation for a shadow with no matching indices.
     The result should be 0 as there are no matching
     """
-    q0, q1, q2 = cirq.LineQubit.range(3)
-    observable = mitiq.PauliString("XYZ", support=(0, 1, 2))
+    pauli = mitiq.PauliString("XYZ")
     measurement_outcomes = ["101", "010", "101"], ["ZXY", "YZX", "ZZY"]
-    k_shadows = 1
+    batch_size = 1
 
     result = expectation_estimation_shadow(
-        measurement_outcomes, observable, k_shadows, False
+        measurement_outcomes, pauli, batch_size
     )
 
-    assert result == 0.0
+    assert result == 0
diff --git a/mitiq/shadows/test/test_shadows.py b/mitiq/shadows/test/test_shadows.py
index 795681ee0d..3701575855 100644
--- a/mitiq/shadows/test/test_shadows.py
+++ b/mitiq/shadows/test/test_shadows.py
@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 """Test classical shadow estimation process."""
+from numbers import Number
 
 import cirq
 
@@ -38,7 +39,6 @@ def executor(
 
 
 def test_pauli_twirling_calibrate():
-
     # Call the function with valid inputs
     result = pauli_twirling_calibrate(
         qubits=qubits, executor=executor, num_total_measurements_calibration=2
@@ -47,9 +47,8 @@ def test_pauli_twirling_calibrate():
     # Check that the dictionary contains the correct number of entries
     assert len(result) <= 2**num_qubits
 
-    # Check that all values in the dictionary are floats
     for value in result.values():
-        assert isinstance(value, complex)
+        assert isinstance(value, Number)
 
     # Call shadow_quantum_processing to get shadow_outcomes
     shadow_outcomes = (["11", "00"], ["ZZ", "XX"])
@@ -63,13 +62,11 @@ def test_pauli_twirling_calibrate():
     # Check that the dictionary contains the correct number of entries
     assert len(result) <= 2**num_qubits
 
-    # Check that all values in the dictionary are floats
     for value in result.values():
-        assert isinstance(value, complex)
+        assert isinstance(value, Number)
 
 
 def test_shadow_quantum_processing():
-
     # Call the function with valid inputs
     result = shadow_quantum_processing(
         circuit, executor, num_total_measurements_shadow=10
@@ -112,7 +109,6 @@ def test_classical_post_processing():
     )
     result_cal = classical_post_processing(
         shadow_outcomes,
-        use_calibration=True,
         calibration_results=calibration_results,
         observables=observables,
         k_shadows=1,
diff --git a/mitiq/shadows/test/test_shadows_utils.py b/mitiq/shadows/test/test_shadows_utils.py
index 459e6b4f6d..b08e19fa00 100644
--- a/mitiq/shadows/test/test_shadows_utils.py
+++ b/mitiq/shadows/test/test_shadows_utils.py
@@ -3,56 +3,54 @@
 # This source code is licensed under the GPL license (v3) found in the
 # LICENSE file in the root directory of this source tree.
 
-"""Defines utility functions for classical shadows protocol."""
+import math
 
 import numpy as np
 
 import mitiq
 from mitiq.shadows.shadows_utils import (
-    bitstring_to_eigenvalues,
+    batch_calibration_data,
     create_string,
-    eigenvalues_to_bitstring,
     fidelity,
     n_measurements_opts_expectation_bound,
     n_measurements_tomography_bound,
+    valid_bitstrings,
 )
 
-# Tests start here
 
+def test_create_string():
+    str_len = 5
+    loc_list = [1, 3]
+    assert create_string(str_len, loc_list) == "01010"
 
-def test_eigenvalues_to_bitstring():
-    values = [-1, 1, 1]
-    assert eigenvalues_to_bitstring(values) == "100"
-    assert bitstring_to_eigenvalues(eigenvalues_to_bitstring(values)) == values
 
+def test_valid_bitstrings():
+    num_qubits = 5
+    bitstrings_on_5_qubits = valid_bitstrings(num_qubits)
+    assert len(bitstrings_on_5_qubits) == 2**num_qubits
+    assert all(b == "0" or b == "1" for b in bitstrings_on_5_qubits.pop())
 
-def test_bitstring_to_eigenvalues():
-    bitstring = "100"
-    np.testing.assert_array_equal(
-        bitstring_to_eigenvalues(bitstring), np.array([-1, 1, 1])
-    )
-    assert (
-        eigenvalues_to_bitstring(bitstring_to_eigenvalues(bitstring))
-        == bitstring
+    num_qubits = 4
+    max_hamming_weight = 2
+    bitstrings_on_3_qubits_hamming_2 = valid_bitstrings(
+        num_qubits, max_hamming_weight
     )
+    assert len(bitstrings_on_3_qubits_hamming_2) == sum(
+        math.comb(num_qubits, i) for i in range(max_hamming_weight + 1)
+    )  # sum_{i == 0}^{max_hamming_weight} (num_qubits choose i)
 
 
-def test_create_string():
-    str_len = 5
-    loc_list = [1, 3]
-    assert create_string(str_len, loc_list) == "01010"
+def test_batch_calibration_data():
+    data = (["010", "110", "000", "001"], ["XXY", "ZYY", "ZZZ", "XYZ"])
+    num_batches = 2
+    for bits, paulis in batch_calibration_data(data, num_batches):
+        assert len(bits) == len(paulis) == num_batches
 
 
 def test_n_measurements_tomography_bound():
-    assert (
-        n_measurements_tomography_bound(0.5, 2) == 2176
-    ), f"Expected 2176, got {n_measurements_tomography_bound(0.5, 2)}"
-    assert (
-        n_measurements_tomography_bound(1.0, 1) == 136
-    ), f"Expected 136, got {n_measurements_tomography_bound(1.0, 1)}"
-    assert (
-        n_measurements_tomography_bound(0.1, 3) == 217599
-    ), f"Expected 217599, got {n_measurements_tomography_bound(0.1, 3)}"
+    assert n_measurements_tomography_bound(0.5, 2) == 2176
+    assert n_measurements_tomography_bound(1.0, 1) == 136
+    assert n_measurements_tomography_bound(0.1, 3) == 217599
 
 
 def test_n_measurements_opts_expectation_bound():
@@ -62,8 +60,8 @@ def test_n_measurements_opts_expectation_bound():
         mitiq.PauliString("Z"),
     ]
     N, K = n_measurements_opts_expectation_bound(0.5, observables, 0.1)
-    assert isinstance(N, int), f"Expected int, got {type(N)}"
-    assert isinstance(K, int), f"Expected int, got {type(K)}"
+    assert isinstance(N, int)
+    assert isinstance(K, int)
 
 
 def test_fidelity():

From ee01c95496e48f2a5f1c5e2fdcc84b7ac50513c2 Mon Sep 17 00:00:00 2001
From: Sam Burdick <samuelmburdick@gmail.com>
Date: Mon, 5 Feb 2024 10:55:13 -0800
Subject: [PATCH 2/9] Fix typo in README.md (#2173)

* Fix typo

* Rerun build
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index ce1f4365a4..639c41baec 100644
--- a/README.md
+++ b/README.md
@@ -70,7 +70,7 @@ Error w Mitiq:    0.073
 ### Calibration
 
 Unsure which error mitigation technique or parameters to use?
-Try out the calibration module demonstrated below to help find the best parameters for you particular backend!
+Try out the calibration module demonstrated below to help find the best parameters for your particular backend!
 
 ![](docs/source/img/calibration.gif)
 

From 570a2898b2c41c70a38413bc833e11adb0e4e3d2 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 5 Feb 2024 10:58:13 -0800
Subject: [PATCH 3/9] Bump pytest from 7.1.3 to 8.0.0 (#2167)

Bumps [pytest](https://github.com/pytest-dev/pytest) from 7.1.3 to 8.0.0.
- [Release notes](https://github.com/pytest-dev/pytest/releases)
- [Changelog](https://github.com/pytest-dev/pytest/blob/main/CHANGELOG.rst)
- [Commits](https://github.com/pytest-dev/pytest/compare/7.1.3...8.0.0)

---
updated-dependencies:
- dependency-name: pytest
  dependency-type: direct:development
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 dev_requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dev_requirements.txt b/dev_requirements.txt
index 6990b49bf8..488e1065c0 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -8,7 +8,7 @@ pennylane~=0.34.0
 amazon-braket-sdk~=1.66.0
 
 # Unit tests, coverage, and formatting/style.
-pytest==7.1.3
+pytest==8.0.0
 pytest-xdist[psutil]==3.0.2
 pytest-cov==4.0.0
 flake8==7.0.0

From 4a2b9ea99930aacb7f10a52990ab4755d3c1f2b7 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 5 Feb 2024 10:59:03 -0800
Subject: [PATCH 4/9] Update qiskit-aer requirement from ~=0.13.1 to ~=0.13.2
 (#2157)

Updates the requirements on [qiskit-aer](https://github.com/Qiskit/qiskit-aer) to permit the latest version.
- [Release notes](https://github.com/Qiskit/qiskit-aer/releases)
- [Commits](https://github.com/Qiskit/qiskit-aer/compare/0.13.1...0.13.2)

---
updated-dependencies:
- dependency-name: qiskit-aer
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 dev_requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dev_requirements.txt b/dev_requirements.txt
index 488e1065c0..d0370f3537 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -1,6 +1,6 @@
 # Third-party integration.
 qiskit~=0.45.1
-qiskit-aer~=0.13.1
+qiskit-aer~=0.13.2
 qiskit-ibm-provider~=0.8.0
 pyquil~=3.5.4
 pennylane-qiskit~=0.34.0

From 1eab631bb315c981c0e865f7e2bebe2229b3a4c4 Mon Sep 17 00:00:00 2001
From: Nathan Shammah <nathan.shammah@gmail.com>
Date: Mon, 5 Feb 2024 20:00:22 +0100
Subject: [PATCH 5/9] Create SECURITY.md (#2162)

* Create SECURITY.md

Create security policy document with basic info.

* remove template
---
 SECURITY.md | 9 +++++++++
 1 file changed, 9 insertions(+)
 create mode 100644 SECURITY.md

diff --git a/SECURITY.md b/SECURITY.md
new file mode 100644
index 0000000000..f84cd98961
--- /dev/null
+++ b/SECURITY.md
@@ -0,0 +1,9 @@
+# Security Policy
+
+## Supported Versions
+
+The latest minor version published on the PyPI server includes the latest security updates.
+
+## Reporting a Vulnerability
+
+Please email us at [info@unitary.fund](mailto:info@unitary.fund) to report a vulnerability or alternatively [open an issue](https://github.com/unitaryfund/mitiq/issues/new). In both cases Unitary Fund will be in touch to triage the issue and respond.

From 2205eb8f5aa620eca03ca60e1d83dd9e53d388fa Mon Sep 17 00:00:00 2001
From: Misty Wahl <82074193+Misty-W@users.noreply.github.com>
Date: Mon, 5 Feb 2024 11:48:15 -0800
Subject: [PATCH 6/9] Reduce doc build time for learning representations
 (#2165)

* load pre-executed pec values from file

* Fix file path for loading PEC data

* Remove duplicate lines

* Fix typo in circuit causing wrong loss fn vals
---
 .../examples/learning-depolarizing-noise.md   | 36 +++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/docs/source/examples/learning-depolarizing-noise.md b/docs/source/examples/learning-depolarizing-noise.md
index 79965bf906..2419e1c3be 100644
--- a/docs/source/examples/learning-depolarizing-noise.md
+++ b/docs/source/examples/learning-depolarizing-noise.md
@@ -50,7 +50,7 @@ Here we use a simple Rx-Rz-CNOT circuit, with an (optional) seed for reproducibi
 
 ```{code-cell} ipython3
 circuit = random_x_z_cnot_circuit(
-    LineQubit.range(2), n_moments=10, random_state=np.random.RandomState(1)
+    LineQubit.range(2), n_moments=5, random_state=np.random.RandomState(1)
 )
 print(circuit)
 ```
@@ -110,7 +110,8 @@ warnings.simplefilter("ignore", UserWarning)
 ```
 
 ```{code-cell} ipython3
-:tags: ["skip-execution"]
+:tags: [skip-execution]
+
 training_circuits = generate_training_circuits(
     circuit=circuit,
     num_training_circuits=5,
@@ -163,6 +164,35 @@ Here we use the Nelder-Mead method in [`scipy.optimize.minimize`](https://docs.s
 to find the value of epsilon that minimizes the depolarizing noise loss function.
 
 ```{code-cell} ipython3
+:tags: [remove-cell]
+
+import os
+
+eps_string = str(epsilon).replace(".", "_")
+pec_data = np.loadtxt(
+        os.path.join(
+            "../../../mitiq/pec/representations/tests/learning_pec_data",
+            f"learning_pec_data_eps_{eps_string}.txt",
+        )
+    )
+
+[success, epsilon_opt] = learn_depolarizing_noise_parameter(
+    operations_to_learn,
+    circuit,
+    ideal_executor,
+    noisy_executor,
+    num_training_circuits=5,
+    fraction_non_clifford=0.2,
+    training_random_state=np.random.RandomState(1),
+    epsilon0=epsilon0,
+    observable=observable,
+    learning_kwargs={"pec_data": pec_data},
+)
+```
+
+```{code-cell} ipython3
+:tags: [skip-execution]
+
 [success, epsilon_opt] = learn_depolarizing_noise_parameter(
     operations_to_learn,
     circuit,
@@ -175,7 +205,9 @@ to find the value of epsilon that minimizes the depolarizing noise loss function
     epsilon0=epsilon0,
     observable=observable,
 )
+```
 
+```{code-cell} ipython3
 print(success)
 print(f"Difference of learned value from true value: {abs(epsilon_opt - epsilon) :.5f}")
 print(f"Difference of initial guess from true value: {abs(epsilon0 - epsilon) :.5f}")

From 6210b25106fb3f12fdf9554940f1fa1598da39d4 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 5 Feb 2024 13:09:02 -0800
Subject: [PATCH 7/9] Update amazon-braket-sdk requirement from ~=1.66.0 to
 ~=1.68.3 (#2175)

Updates the requirements on [amazon-braket-sdk](https://github.com/amazon-braket/amazon-braket-sdk-python) to permit the latest version.
- [Release notes](https://github.com/amazon-braket/amazon-braket-sdk-python/releases)
- [Changelog](https://github.com/amazon-braket/amazon-braket-sdk-python/blob/main/CHANGELOG.md)
- [Commits](https://github.com/amazon-braket/amazon-braket-sdk-python/compare/v1.66.0...v1.68.3)

---
updated-dependencies:
- dependency-name: amazon-braket-sdk
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 dev_requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dev_requirements.txt b/dev_requirements.txt
index d0370f3537..025ba8ba0c 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -5,7 +5,7 @@ qiskit-ibm-provider~=0.8.0
 pyquil~=3.5.4
 pennylane-qiskit~=0.34.0
 pennylane~=0.34.0
-amazon-braket-sdk~=1.66.0
+amazon-braket-sdk~=1.68.3
 
 # Unit tests, coverage, and formatting/style.
 pytest==8.0.0

From a2f3450cd2c492781711b3011c000927616db9b7 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Wed, 7 Feb 2024 17:16:39 -0800
Subject: [PATCH 8/9] Bump pyscf from 2.4.0 to 2.5.0 (#2176)

Bumps [pyscf](https://github.com/pyscf/pyscf) from 2.4.0 to 2.5.0.
- [Release notes](https://github.com/pyscf/pyscf/releases)
- [Changelog](https://github.com/pyscf/pyscf/blob/master/CHANGELOG)
- [Commits](https://github.com/pyscf/pyscf/compare/v2.4.0...v2.5.0)

---
updated-dependencies:
- dependency-name: pyscf
  dependency-type: direct:development
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 dev_requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dev_requirements.txt b/dev_requirements.txt
index 025ba8ba0c..724429ff80 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -29,7 +29,7 @@ sphinx-gallery==0.10.1
 nbsphinx==0.9.1
 matplotlib==3.8.1
 pandas==2.1.3
-pyscf==2.4.0; sys_platform != 'win32'
+pyscf==2.5.0; sys_platform != 'win32'
 openfermion==1.6.0; sys_platform != 'win32'
 openfermionpyscf==0.5; sys_platform != 'win32'
 bqskit[ext]==1.0.4

From 70a325e2ba529c333902dda3e81aada7ee42a845 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Wed, 7 Feb 2024 17:20:57 -0800
Subject: [PATCH 9/9] Update amazon-braket-sdk requirement from ~=1.68.3 to
 ~=1.69.0 (#2177)

Updates the requirements on [amazon-braket-sdk](https://github.com/amazon-braket/amazon-braket-sdk-python) to permit the latest version.
- [Release notes](https://github.com/amazon-braket/amazon-braket-sdk-python/releases)
- [Changelog](https://github.com/amazon-braket/amazon-braket-sdk-python/blob/main/CHANGELOG.md)
- [Commits](https://github.com/amazon-braket/amazon-braket-sdk-python/compare/v1.68.3...v1.69.0)

---
updated-dependencies:
- dependency-name: amazon-braket-sdk
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
---
 dev_requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dev_requirements.txt b/dev_requirements.txt
index 724429ff80..7a38e1923c 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -5,7 +5,7 @@ qiskit-ibm-provider~=0.8.0
 pyquil~=3.5.4
 pennylane-qiskit~=0.34.0
 pennylane~=0.34.0
-amazon-braket-sdk~=1.68.3
+amazon-braket-sdk~=1.69.0
 
 # Unit tests, coverage, and formatting/style.
 pytest==8.0.0