More debugging for edge cases and code coverage

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()

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():