diff --git a/qiskit_addon_utils/exp_vals/expectation_values.py b/qiskit_addon_utils/exp_vals/expectation_values.py index 281f4f7..005e505 100644 --- a/qiskit_addon_utils/exp_vals/expectation_values.py +++ b/qiskit_addon_utils/exp_vals/expectation_values.py @@ -109,6 +109,8 @@ def executor_expectation_values( postselect_mask = postselect_mask.reshape((1, *postselect_mask.shape)) meas_basis_axis = 0 avg_axis = tuple(a + 1 for a in avg_axis) + elif meas_basis_axis < 0: + raise ValueError("meas_basis_axis must be nonnegative.") if len(basis_dict) != bool_array.shape[meas_basis_axis]: raise ValueError( @@ -157,13 +159,30 @@ def executor_expectation_values( ##### PEC SIGNS: if pauli_signs is not None: bool_array, basis_dict, net_signs = _apply_pec_signs(bool_array, basis_dict, pauli_signs) - else: - net_signs = np.zeros(bool_array.shape[:-2], dtype=bool) - # For PEC, we will need to apply a rescaling factor gamma later when computing expectation values. + # For PEC, we will need to multiply by gamma later when computing expectation values. + if gamma_factor is None: + # If gamma not provided, estimate it empirically, for each requested expectation value: + gamma_factor = 1 / (1 - 2 * np.mean(net_signs, axis=avg_axis)) + elif gamma_factor is None: + gamma_factor = 1.0 + + ##### If other axes are to be averaged over, do so by first absorbing them into the shots axis: + if avg_axis: + # move avg_axis just before shots axis (just before axis -2): + axis_positions_before_shots = -2 - np.arange(len(avg_axis)) + bool_array = np.moveaxis(bool_array, avg_axis, axis_positions_before_shots) + # flatten into shots axis (preserve sizes of other axes, including bits axis): + bool_array = np.reshape( + bool_array, (*bool_array.shape[: -2 - len(avg_axis)], -1, bool_array.shape[-1]) + ) + + # update others to match: + num_shots_kept = np.sum(num_shots_kept, avg_axis) + meas_basis_axis -= int(np.sum(np.asarray(avg_axis) < meas_basis_axis)) ##### ACCUMULATE CONTRIBUTIONS FROM EACH MEAS BASIS: barray = BitArray.from_bool_array(bool_array, "little") - output_shape_each_obs = np.delete(barray.shape, (meas_basis_axis, *avg_axis)) + output_shape_each_obs = np.delete(barray.shape, meas_basis_axis) mean_each_observable = np.zeros((num_observables, *output_shape_each_obs), dtype=float) var_each_observable = np.zeros((num_observables, *output_shape_each_obs), dtype=float) @@ -175,7 +194,6 @@ def executor_expectation_values( idx = tuple([*skip_axes, meas_basis_idx]) barray_this_basis = barray[idx] num_kept = num_shots_kept[idx] - signs = net_signs[idx] basis_rescale_factors = ( rescale_factors[meas_basis_idx] if rescale_factors is not None else None ) @@ -188,32 +206,12 @@ def executor_expectation_values( rescale_each_observable=basis_rescale_factors, ) - variances = standard_errs**2 - del standard_errs - - ## AVERAGE OVER SPECIFIED AXES ("TWIRLS"): - # Update indexing since we already sliced away meas_basis axis: - avg_axis_ = tuple(a if a < meas_basis_axis else a - 1 for a in avg_axis) - - if gamma_factor is not None: - axes = tuple(ax for i, ax in enumerate(signs.shape) if i not in avg_axis_) - rescaling = np.full(axes, gamma_factor) - else: - num_minus = np.count_nonzero(signs, axis=avg_axis_) - num_plus = np.count_nonzero(~signs, axis=avg_axis_) - num_twirls = num_plus + num_minus - rescaling = num_twirls / (num_plus - num_minus) - - # Will weight each twirl by its fraction of kept shots. - # If no postselection, weighting reduces to dividing by num_twirls: - num_kept_each_avg = np.sum(num_kept, axis=avg_axis_) - weights = num_kept / np.expand_dims(num_kept_each_avg, avg_axis_) # Append axis for observables being evaluated (to match axis in `means`): - weights = weights[..., np.newaxis] - rescaling = rescaling[..., np.newaxis] - means = rescaling * np.sum(means * weights, axis=avg_axis_) + if not isinstance(gamma_factor, float): + gamma_factor = gamma_factor[..., np.newaxis] + means = gamma_factor * means # Propagate uncertainties: - variances = rescaling**2 * np.sum(variances * weights**2, axis=avg_axis_) + variances = (gamma_factor * standard_errs) ** 2 # Move observable axis from end to front: mean_each_observable += np.moveaxis(means, -1, 0) var_each_observable += np.moveaxis(variances, -1, 0) @@ -386,16 +384,15 @@ def _bitarray_expectation_value( parities = (outcomes & mask_z).bitcount() % 2 # Compute the coefficients of 0 and 1 components. - nulled_by_0_projector = np.logical_or.reduce((outcomes & mask_0).array, axis=-1) - nulled_by_1_projector = np.logical_or.reduce(((outcomes & mask_1) ^ mask_1).array, axis=-1) - coeffs_01 = ~np.logical_or(nulled_by_0_projector, nulled_by_1_projector) + nulled_by_projector = np.any((outcomes & mask_0).array, axis=-1) + nulled_by_projector |= np.any((~outcomes & mask_1).array, axis=-1) + coeffs_01 = ~nulled_by_projector # Compute expectation values shape = np.broadcast_shapes(outcomes.shape, mask_z.shape) expvals_each_term = np.zeros(shape, dtype=float) sq_expvals_each_term = np.zeros_like(expvals_each_term) - # Combine masks to get coeff for each shot (-1, 0, or 1) if np.all(coeffs_01): # We can do a faster computation of pure Pauli parities expvals_each_term += outcomes.num_shots - 2 * np.sum(parities, axis=-1, dtype=int) @@ -407,15 +404,15 @@ def _bitarray_expectation_value( # Divide by total shots. May be less than nominal number in array if # we are postselecting via projector in observable terms: - shots = np.asarray(outcomes.num_shots) if shots is None else shots[..., np.newaxis] + shots = outcomes.num_shots if shots is None else shots[..., np.newaxis] # Edge case of counts dict containing outcomes but with total shots, eg {"0": 0}. with np.errstate(divide="ignore"): expvals_each_term /= shots sq_expvals_each_term /= shots - expvals_each_term[~np.isfinite(expvals_each_term)] = np.nan - sq_expvals_each_term[~np.isfinite(expvals_each_term)] = np.nan variances_each_term = np.clip(sq_expvals_each_term - expvals_each_term**2, 0, None) / shots + expvals_each_term[~np.isfinite(expvals_each_term)] = np.nan + variances_each_term[~np.isfinite(variances_each_term)] = np.nan # len(all_coeffs) == number of bit terms == expvals_each_term.shape[-1], so broadcasts: # TODO: test case of empty observable