diff --git a/src/qibochem/measurement/optimization.py b/src/qibochem/measurement/optimization.py index 1320b15..8defbd1 100644 --- a/src/qibochem/measurement/optimization.py +++ b/src/qibochem/measurement/optimization.py @@ -72,7 +72,7 @@ def allocate_shots(grouped_terms, n_shots, method=None, max_shots_per_term=None) term_coefficients = np.array( [sum(abs(term.coefficient.real) for term in terms) for (_, terms) in grouped_terms] ) - max_shots_per_term = int(n_shots * (np.max(term_coefficients) / sum(term_coefficients))) + max_shots_per_term = int(np.ceil(n_shots * (np.max(term_coefficients) / sum(term_coefficients)))) max_shots_per_term = min(max_shots_per_term, 250) # Can be explored further? # Don't let max_shots_per_term exceed the total number of shots max_shots_per_term = min(n_shots, max_shots_per_term) @@ -90,6 +90,7 @@ def allocate_shots(grouped_terms, n_shots, method=None, max_shots_per_term=None) if np.min(shot_allocation) == max_shots_per_term: max_shots_per_term = min(2 * max_shots_per_term, n_shots) shot_allocation = np.zeros(n_terms, dtype=int) + continue if method in ("c", "coefficients"): # Split shots based on the relative magnitudes of the coefficients of the (group of) Pauli term(s) @@ -127,8 +128,7 @@ def allocate_shots(grouped_terms, n_shots, method=None, max_shots_per_term=None) raise NameError("Unknown method!") # Add on to the current allocation, and set upper limit to the number of shots for a given term - shot_allocation += _shot_allocation + shot_allocation += _shot_allocation.astype(int) shot_allocation = np.clip(shot_allocation, 0, max_shots_per_term) return shot_allocation.tolist() - # return shot_allocation.astype(int).tolist() diff --git a/tests/test_expectation_samples.py b/tests/test_expectation_samples.py index b42d612..c779df1 100644 --- a/tests/test_expectation_samples.py +++ b/tests/test_expectation_samples.py @@ -87,9 +87,14 @@ def test_allocate_shots_coefficient(): 190, 10, ], "max_shots_per_term error: Too big test" - # Too few shots to allocate - n_shots = 2 - assert allocate_shots(grouped_terms, n_shots=n_shots) == [1, 1] + + +def test_allocate_shots_coefficient_edges(): + """Edge cases of allocate_shots""" + hamiltonian = SymbolicHamiltonian(Z(0) + X(0)) + grouped_terms = measurement_basis_rotations(hamiltonian, 1) + n_shots = 1 + assert allocate_shots(grouped_terms, n_shots=n_shots) in ([0, 1], [1, 0]) def test_allocate_shots_input_validity():