Simplify expectation_estimation_shadow code

[andreamari]

In the function expectation_estimation_shadow the following code:

        # assign the splits temporarily
        b_lists_k, u_lists_k = (
            b_lists[i : i + n_total_measurements // k_shadows],
            u_lists[i : i + n_total_measurements // k_shadows],
        )
        # number of measurements/shadows in each split
        n_group_measurements = len(b_lists_k)
        # find the exact matches for the observable of

       # interest at the specified locations
        indices = np.all(u_lists_k[:, target_locs] == target_obs, axis=1)

        # catch the edge case where there is no match in the chunk
        if sum(indices) > 0:
            # take the product and sum
            product = np.prod(
                3.0 * (b_lists_k[indices][:, target_locs]), axis=1
            )
            means.append(np.sum(product) / n_group_measurements)
        else:
            means.append(0.0)

is equivalent to the following simplified code:

      mean_k = 0.0
      for b_list, u_list in zip(b_lists_k, u_lists_k):
          if np.all(u_list[target_locs] == target_obs):
              mean_k += 3.0 * np.prod(b_list[target_locs])
      mean.append(mean_k)

The simplified code may be slower since Python for loops are slow. The scope of this issue is to verify if the speed reduction is significant. E.g. by timing the execution of classical shadow test files or of a classical shadow notebook/example. If the performance reduction is significant, please ignore and close this issue. If the performance reduction is small, please replace the long code with the simplified code.

[tinaoberoi]

Hi @andreamari I would like to work on this issue if its still active.

[andreamari]

Thanks, @tinaoberoi! Yes, the issue is still active.

[tinaoberoi]

@andreamari The code for expectation_estimation_shadow has been modified by PR to

        b_lists_shadow_k = b_lists_shadow[idxes]
        u_lists_shadow_k = u_lists_shadow[idxes]
        # number of measurements/shadows in each split
        n_group_measurements = len(b_lists_shadow_k)

        indices = np.all(
            u_lists_shadow_k[:, target_locs] == target_obs, axis=1
        )
        if sum(indices) == 0:
            means.append(0.0)
        else:
            eigenvalues = np.array(
                [
                    bitstring_to_eigenvalues(b)
                    for b in b_lists_shadow_k[indices]
                ]
            )
            product = np.prod(eigenvalues[:, target_locs], 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:
                    means.append(0.0)
                else:
                    # product becomes an array of snapshots expectation values
                    # witch satisfy condition (1) and (2)
                    product = (1.0 / f_val) * product
            else:
                product = 3 ** (target_support.count("1")) * product

            # append the mean of the product in each split
            means.append(np.sum(product) / n_group_measurements)

I guess the simplified code will also have to change for the same.

[andreamari]

Good point! The simplified code must be equivalent of the current code that you posted. The code snippets in the description of this issue are not up to date, but they still give the idea of a possible way to simplify the code.

[bdg221]

If this issue has gone quiet, I can take it over.